edit respond

This commit is contained in:
2025-08-04 12:59:11 +07:00
parent 82b7be53e6
commit 6ec3b0ed04

View File

@@ -1,8 +1,8 @@
from typing import Any, Dict, List, Optional, AsyncGenerator from typing import Any, Dict, List, Optional, AsyncGenerator
import gradio as gr import gradio as gr
import uuid
import torch import torch
import asyncio import asyncio
import threading
from transformers import AutoTokenizer as HFTokenizer from transformers import AutoTokenizer as HFTokenizer
from loguru import logger from loguru import logger
from langchain_chroma import Chroma from langchain_chroma import Chroma
@@ -36,6 +36,7 @@ class ChatWithAI:
max_model_len=32768, max_model_len=32768,
enable_prefix_caching=True, enable_prefix_caching=True,
tensor_parallel_size=1, # Увеличь, если у тебя несколько GPU tensor_parallel_size=1, # Увеличь, если у тебя несколько GPU
enable_chunked_prefill=False,
) )
self.engine = AsyncLLMEngine.from_engine_args(engine_args) self.engine = AsyncLLMEngine.from_engine_args(engine_args)
self.tokenizer = HFTokenizer.from_pretrained(settings.LM_MODEL_NAME, trust_remote_code=True) self.tokenizer = HFTokenizer.from_pretrained(settings.LM_MODEL_NAME, trust_remote_code=True)
@@ -130,7 +131,7 @@ class ChatWithAI:
buffer = "" buffer = ""
# Генерируем асинхронно # Генерируем асинхронно
request_id = f"request-{hash(query)}" request_id = str(uuid.uuid4())
try: try:
async for output in self.engine.generate(prompt, sampling_params, request_id): async for output in self.engine.generate(prompt, sampling_params, request_id):
if output.outputs: if output.outputs:
@@ -154,27 +155,14 @@ class ChatWithAI:
def main(): def main():
chat = ChatWithAI(provider="qwen3") chat = ChatWithAI(provider="qwen3")
# Gradio требует синхронную функцию, но мы можем обернуть асинхронный вызов async def respond(message, history):
def respond(message, history): # Генерируем ответ по частям
# Запускаем асинхронный генератор в синхронном контексте try:
async def async_generate():
async for token in chat.generate_response_stream(message): async for token in chat.generate_response_stream(message):
yield token yield token
except Exception as e:
# Используем asyncio.new_event_loop() для запуска внутри потока logger.error(f"Ошибка в respond: {e}")
loop = asyncio.new_event_loop() yield "Извините, произошла ошибка при обработке запроса."
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()
demo = gr.ChatInterface( demo = gr.ChatInterface(
fn=respond, fn=respond,