edit respond
This commit is contained in:
30
app/main.py
30
app/main.py
@@ -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,
|
||||||
|
|||||||
Reference in New Issue
Block a user