diff --git a/app/main.py b/app/main.py index 4c14656..3432cd8 100644 --- a/app/main.py +++ b/app/main.py @@ -1,8 +1,8 @@ from typing import Any, Dict, List, Optional, AsyncGenerator import gradio as gr +import uuid import torch import asyncio -import threading from transformers import AutoTokenizer as HFTokenizer from loguru import logger from langchain_chroma import Chroma @@ -36,6 +36,7 @@ class ChatWithAI: max_model_len=32768, enable_prefix_caching=True, tensor_parallel_size=1, # Увеличь, если у тебя несколько GPU + enable_chunked_prefill=False, ) self.engine = AsyncLLMEngine.from_engine_args(engine_args) self.tokenizer = HFTokenizer.from_pretrained(settings.LM_MODEL_NAME, trust_remote_code=True) @@ -130,7 +131,7 @@ class ChatWithAI: buffer = "" # Генерируем асинхронно - request_id = f"request-{hash(query)}" + request_id = str(uuid.uuid4()) try: async for output in self.engine.generate(prompt, sampling_params, request_id): if output.outputs: @@ -154,27 +155,14 @@ class ChatWithAI: def main(): chat = ChatWithAI(provider="qwen3") - # Gradio требует синхронную функцию, но мы можем обернуть асинхронный вызов - def respond(message, history): - # Запускаем асинхронный генератор в синхронном контексте - async def async_generate(): + async def respond(message, history): + # Генерируем ответ по частям + try: async for token in chat.generate_response_stream(message): yield token - - # Используем asyncio.new_event_loop() для запуска внутри потока - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - - try: - gen = async_generate() - while True: - try: - token = loop.run_until_complete(gen.__anext__()) - yield token - except StopAsyncIteration: - break - finally: - loop.close() + except Exception as e: + logger.error(f"Ошибка в respond: {e}") + yield "Извините, произошла ошибка при обработке запроса." demo = gr.ChatInterface( fn=respond, @@ -193,4 +181,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main()