Update app.py

app.py

@@ -2,10 +2,9 @@ import os
 import torch
 import gradio as gr
 import requests
-from typing import List, Dict
-from threading import Lock
-
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from typing import List, Dict, Iterator
+from threading import Thread
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextIteratorStreamer
 from peft import PeftModel
 
 # --- 1. Конфигурация и загрузка модели ---
@@ -24,7 +23,7 @@ TOP_P = 0.9
 REPETITION_PENALTY = 1.05
 SYS_PROMPT_TT = (
     "Син - татар цифрлы ярдәмчесе. Татар телендә һәрвакыт ачык һәм дустанә җавап бир."
-    "мәгълүмат җитәрлек булмаса, 1-2 кыска аныклаучы сорау бир. "
+    "мәгълүmat җитәрлек булмаса, 1-2 кыска аныклаучы сорау бир. "
     "Һәрвакыт татарча гына җавап бир."
 )
 
@@ -40,13 +39,11 @@ model.config.use_cache = True
 model.eval()
 print("✅ Модель успешно загружена!")
 
-# --- 2. Логика приложения (функции перевода и генерации) ---
+# --- 2. Логика приложения (с изменениями для стриминга) ---
 
 YANDEX_TRANSLATE_URL = "https://translate.api.cloud.yandex.net/translate/v2/translate"
-YANDEX_DETECT_URL = "https://translate.api.cloud.yandex.net/translate/v2/detect" # НОВЫЙ URL ДЛЯ ОПРЕДЕЛЕНИЯ ЯЗЫКА
-generation_lock = Lock()
+YANDEX_DETECT_URL = "https://translate.api.cloud.yandex.net/translate/v2/detect"
 
-# НОВАЯ ФУНКЦИЯ для определения языка 🧠
 def detect_language(text: str) -> str:
     headers = {"Authorization": f"Api-Key {YANDEX_API_KEY}"}
     payload = {"folderId": YANDEX_FOLDER_ID, "text": text}
@@ -54,40 +51,32 @@ def detect_language(text: str) -> str:
         resp = requests.post(YANDEX_DETECT_URL, headers=headers, json=payload, timeout=10)
         resp.raise_for_status()
         data = resp.json()
-        return data.get("languageCode", "ru")  # Если не определился, считаем, что русский
+        return data.get("languageCode", "ru")
     except requests.exceptions.RequestException as e:
         print(f"Ошибка определения языка: {e}")
-        return "ru" # В случае ошибки считаем, что это русский для безопасности
+        return "ru"
 
-def _yandex_translate(texts: List[str], source: str, target: str) -> List[str]:
+def ru2tt(text: str) -> str:
     headers = {"Authorization": f"Api-Key {YANDEX_API_KEY}"}
-    payload = {"folderId": YANDEX_FOLDER_ID, "texts": texts, "sourceLanguageCode": source, "targetLanguageCode": target}
+    payload = {"folderId": YANDEX_FOLDER_ID, "texts": [text], "sourceLanguageCode": "ru", "targetLanguageCode": "tt"}
     try:
         resp = requests.post(YANDEX_TRANSLATE_URL, headers=headers, json=payload, timeout=30)
         resp.raise_for_status()
-        data = resp.json()
-        return [item["text"] for item in data["translations"]]
+        return resp.json()["translations"][0]["text"]
     except requests.exceptions.RequestException as e:
         print(f"Ошибка перевода: {e}")
-        return [f"Ошибка перевода: {text}" for text in texts]
-
-def ru2tt(text: str) -> str:
-    return _yandex_translate([text], "ru", "tt")[0]
-
-def tt2ru(text: str) -> str:
-    return _yandex_translate([text], "tt", "ru")[0]
+        return f"Ошибка перевода: {text}"
 
 def render_prompt(messages: List[Dict[str, str]]) -> str:
+    # Ваша функция render_prompt остается без изменений
     if getattr(tok, "chat_template", None):
         try:
             return tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        except Exception:
-            pass
+        except Exception: pass
     sys_text = ""
     turns = []
     for m in messages:
-        if m["role"] == "system":
-            sys_text += m["content"].strip() + "\n"
+        if m["role"] == "system": sys_text += m["content"].strip() + "\n"
     i = 0
     while i < len(messages):
         m = messages[i]
@@ -95,32 +84,51 @@ def render_prompt(messages: List[Dict[str, str]]) -> str:
             next_assistant = None
             if i + 1 < len(messages) and messages[i + 1]["role"] == "assistant":
                 next_assistant = messages[i + 1]["content"]
-            if len(turns) == 0 and sys_text:
-                user_block = f"<<SYS>>\n{sys_text.strip()}\n<</SYS>>\n\n{m['content']}"
-            else:
-                user_block = m["content"]
+            user_block = f"<<SYS>>\n{sys_text.strip()}\n<</SYS>>\n\n{m['content']}" if len(turns) == 0 and sys_text else m['content']
             if next_assistant is None:
                 turns.append(f"<s>[INST] {user_block} [/INST]")
             else:
                 turns.append(f"<s>[INST] {user_block} [/INST] {next_assistant}</s>")
                 i += 1
         i += 1
-    if not turns:
-        return f"<s>[INST] <<SYS>>\n{sys_text.strip()}\n<</SYS>>\n\n [/INST]" if sys_text else "<s>[INST] [/INST]"
-    return "".join(turns)
+    return "".join(turns) if turns else (f"<s>[INST] <<SYS>>\n{sys_text.strip()}\n<</SYS>>\n\n [/INST]" if sys_text else "<s>[INST] [/INST]")
 
+# ❗ ИЗМЕНЕННАЯ ФУНКЦИЯ ГЕНЕРАЦИИ
 @torch.inference_mode()
-def generate_tt_reply(messages: List[Dict[str, str]]) -> str:
-    with generation_lock:
-        prompt = render_prompt(messages)
-        inputs = tok(prompt, return_tensors="pt").to(model.device)
-        out = model.generate(**inputs, max_new_tokens=MAX_NEW_TOKENS, do_sample=True, temperature=TEMPERATURE, top_p=TOP_P, repetition_penalty=REPETITION_PENALTY, eos_token_id=tok.eos_token_id, pad_token_id=tok.pad_token_id)
-        gen_text = tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
-        return gen_text.strip()
-
-# --- 3. Gradio интерфейс ---
-
-def chat_fn(message, history):
+def generate_tt_reply_stream(messages: List[Dict[str, str]]) -> Iterator[str]:
+    prompt = render_prompt(messages)
+    inputs = tok(prompt, return_tensors="pt").to(model.device)
+    
+    # Создаем streamer
+    streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True)
+
+    # Аргументы для генерации
+    generation_kwargs = dict(
+        inputs,
+        streamer=streamer,
+        max_new_tokens=MAX_NEW_TOKENS,
+        do_sample=True,
+        temperature=TEMPERATURE,
+        top_p=TOP_P,
+        repetition_penalty=REPETITION_PENALTY,
+        eos_token_id=tok.eos_token_id,
+        pad_token_id=tok.pad_token_id,
+    )
+    
+    # Запускаем генерацию в отдельном потоке
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    # Yield'им каждый новый кусочек текста
+    generated_text = ""
+    for new_text in streamer:
+        generated_text += new_text
+        yield generated_text
+
+# --- 3. Gradio интерфейс (с изменениями для стриминга) ---
+
+# ❗ ИЗМЕНЕННАЯ ФУНКЦИЯ-КОНТРОЛЛЕР
+def chat_fn(message: str, history: list) -> Iterator[list]:
     # 1. Формируем историю для модели
     messages = [{"role": "system", "content": SYS_PROMPT_TT}]
     for user_msg, bot_msg in history:
@@ -128,29 +136,25 @@ def chat_fn(message, history):
         if bot_msg:
              messages.append({"role": "assistant", "content": bot_msg})
 
-    # 2. ОПРЕДЕЛЯЕМ ЯЗЫК и переводим, если нужно 🛡️
+    # 2. Определяем язык и переводим, если нужно
     detected_lang = detect_language(message)
-    if detected_lang != "tt":
-        user_tt = ru2tt(message)
-    else:
-        user_tt = message  # Уже на татарском, используем как есть
-
+    user_tt = ru2tt(message) if detected_lang != "tt" else message
+    
     messages.append({"role": "user", "content": user_tt})
 
-    # 3. Генерируем ответ модели
-    tt_reply = generate_tt_reply(messages)
-
-    # 4. Добавляем в историю татарский вопрос и татарский ответ
-    history.append([user_tt, tt_reply])
+    # 3. Добавляем в историю сообщение пользователя и пустой ответ бота
+    history.append([user_tt, ""])
 
-    # 5. Возвращаем полную историю на татарском
-    return history
+    # 4. Стримим ответ модели и обновляем историю на лету
+    for partial_response in generate_tt_reply_stream(messages):
+        history[-1][1] = partial_response  # Обновляем последнее сообщение в истории
+        yield history # Возвращаем всю историю на каждом шаге
 
-# Создаем интерфейс с татарскими надписями
+# Создаем интерфейс
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("## Татарский чат-бот от команды Сбера")
-    chatbot = gr.Chatbot(label="Диалог", height=500)
-    msg = gr.Textbox(label="Хәбәрегезне рус яки татар телендә языгыз", placeholder="Татарстанның башкаласы нинди шәһәр?")
+    chatbot = gr.Chatbot(label="Диалог", height=500, bubble_full_width=False)
+    msg = gr.Textbox(label="Хәбәрегезне рус яки татар телендә языгыз", placeholder="Татарстанның башкаласы нинди шәһәр? / Какая столица Татарстана?")
     clear = gr.Button("🗑️ Чистарту")
 
     msg.submit(chat_fn, inputs=[msg, chatbot], outputs=chatbot)