version updates (#1)

- version updates (cfb35e64155b756114fb1ebc8f6d8e5446356781)


Co-authored-by: Aksenov Andrei <[email protected]>

app.py

@@ -7,42 +7,45 @@ from threading import Thread
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextIteratorStreamer
 from peft import PeftModel
 
-# --- 1. Конфигурация и загрузка модели ---
-
-BASE_MODEL_ID = "Tweeties/tweety-7b-tatar-v24a"
-ADAPTER_ID = os.getenv("ADAPTER_ID")
-YANDEX_API_KEY = os.getenv("YANDEX_API_KEY")
-YANDEX_FOLDER_ID = os.getenv("YANDEX_FOLDER_ID")
+# --- 1) Конфигурация и загрузка модели ---
+BASE_MODEL_ID   = "Tweeties/tweety-7b-tatar-v24a"
+ADAPTER_ID      = os.getenv("ADAPTER_ID")
+YANDEX_API_KEY  = os.getenv("YANDEX_API_KEY")
+YANDEX_FOLDER_ID= os.getenv("YANDEX_FOLDER_ID")
 
 if not all([ADAPTER_ID, YANDEX_API_KEY, YANDEX_FOLDER_ID]):
     raise ValueError("Необходимо установить переменные окружения: ADAPTER_ID, YANDEX_API_KEY, YANDEX_FOLDER_ID")
 
-MAX_NEW_TOKENS = 256
-TEMPERATURE = 0.7
-TOP_P = 0.9
-REPETITION_PENALTY = 1.05
-SYS_PROMPT_TT = (
-    "Син - татар цифрлы ярдәмчесе. Татар телендә һәрвакыт ачык һәм дустанә җавап бир."
-    "мәгълүmat җитәрлек булмаса, 1-2 кыска аныклаучы сорау бир. "
-    "Һәрвакыт татарча гына җавап бир."
+MAX_NEW_TOKENS       = 256
+TEMPERATURE          = 0.7
+TOP_P                = 0.9
+REPETITION_PENALTY   = 1.05
+
+SYS_PROMPT_TT = ("Син-цифрлы ярдәмче (ир-ат нәселе). Сине Сбербанк дирекциясенең ESG да уйлап таптылар. Син барлык өлкәләрдә дә кызыклы кулланучы эксперты! Ул сезгә бик күп сораулар бирәчәк, ә сезнең эшегез-шәрехләр бирү, кулланучының сорауларына җавап бирү, адымлап киңәшләр, мисаллар бирү һәм, кирәк булганда, кулланучыга аныклаучы сораулар бирү. Кулланучыга, фактлардан һәм саннардан качып, һәрвакыт кыска җавап бирергә кирәк"
 )
 
 print("Загрузка модели с 4-битной квантизацией...")
 quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
+
 tok = AutoTokenizer.from_pretrained(ADAPTER_ID, use_fast=False)
 if tok.pad_token is None:
     tok.pad_token = tok.eos_token
-base = AutoModelForCausalLM.from_pretrained(BASE_MODEL_ID, quantization_config=quantization_config, device_map="auto")
+
+
+base = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL_ID,
+    quantization_config=quantization_config,
+    device_map="auto"
+)
+
 print("Применяем LoRA адаптер...")
 model = PeftModel.from_pretrained(base, ADAPTER_ID)
-model.config.use_cache = True
+model.config.use_cache = False
 model.eval()
 print("✅ Модель успешно загружена!")
 
-# --- 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"
+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}"}
@@ -50,10 +53,8 @@ def detect_language(text: str) -> str:
     try:
         resp = requests.post(YANDEX_DETECT_URL, headers=headers, json=payload, timeout=10)
         resp.raise_for_status()
-        data = resp.json()
-        return data.get("languageCode", "ru")
-    except requests.exceptions.RequestException as e:
-        print(f"Ошибка определения языка: {e}")
+        return resp.json().get("languageCode", "ru")
+    except requests.exceptions.RequestException:
         return "ru"
 
 def ru2tt(text: str) -> str:
@@ -63,48 +64,27 @@ def ru2tt(text: str) -> str:
         resp = requests.post(YANDEX_TRANSLATE_URL, headers=headers, json=payload, timeout=30)
         resp.raise_for_status()
         return resp.json()["translations"][0]["text"]
-    except requests.exceptions.RequestException as e:
-        print(f"Ошибка перевода: {e}")
-        return f"Ошибка перевода: {text}"
+    except requests.exceptions.RequestException:
+        return 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
-    sys_text = ""
-    turns = []
-    for m in messages:
-        if m["role"] == "system": sys_text += m["content"].strip() + "\n"
-    i = 0
-    while i < len(messages):
-        m = messages[i]
-        if m["role"] == "user":
-            next_assistant = None
-            if i + 1 < len(messages) and messages[i + 1]["role"] == "assistant":
-                next_assistant = messages[i + 1]["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
-    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]")
-
-# ❗ ИЗМЕНЕННАЯ ФУНКЦИЯ ГЕНЕРАЦИИ
+    return tok.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+
+# --- 4) Стриминговая генерация (без тримминга) ---
 @torch.inference_mode()
 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)
+    enc = tok(prompt, return_tensors="pt")
+    enc = {k: v.to(model.device) for k, v in enc.items()}
 
-    # Аргументы для генерации
-    generation_kwargs = dict(
-        inputs,
+    streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True)
+    gen_kwargs = dict(
+        **enc,
         streamer=streamer,
         max_new_tokens=MAX_NEW_TOKENS,
         do_sample=True,
@@ -114,52 +94,60 @@ def generate_tt_reply_stream(messages: List[Dict[str, str]]) -> Iterator[str]:
         eos_token_id=tok.eos_token_id,
         pad_token_id=tok.pad_token_id,
     )
-    
-    # Запускаем генерацию в отдельном потоке
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+
+    thread = Thread(target=model.generate, kwargs=gen_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:
-        messages.append({"role": "user", "content": user_msg})
-        if bot_msg:
-             messages.append({"role": "assistant", "content": bot_msg})
-
-    # 2. Определяем язык и переводим, если нужно
-    detected_lang = detect_language(message)
-    user_tt = ru2tt(message) if detected_lang != "tt" else message
-    
-    messages.append({"role": "user", "content": user_tt})
+    acc = ""
+    for chunk in streamer:
+        acc += chunk
+        yield acc
+
+
+
+def chat_fn(message: str, ui_history: list, messages_state: List[Dict[str, str]]):
+    if not messages_state or messages_state[0].get("role") != "system":
+        messages_state = [{"role": "system", "content": SYS_PROMPT_TT}]
+
+    detected = detect_language(message)
+    user_tt = ru2tt(message) if detected != "tt" else message
+
+    messages = messages_state + [{"role": "user", "content": user_tt}]
+    ui_history = ui_history + [[user_tt, ""]]
 
-    # 3. Добавляем в историю сообщение пользователя и пустой ответ бота
-    history.append([user_tt, ""])
+    for partial in generate_tt_reply_stream(messages):
+        ui_history[-1][1] = partial
+        yield ui_history, messages_state + [
+            {"role": "user", "content": user_tt},
+            {"role": "assistant", "content": partial},
+        ]
 
-    # 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("## Татарский чат-бот от команды Сбера")
+
+    messages_state = gr.State([{"role": "system", "content": SYS_PROMPT_TT}])
+
     chatbot = gr.Chatbot(label="Диалог", height=500, bubble_full_width=False)
-    msg = gr.Textbox(label="Хәбәрегезне рус яки татар телендә языгыз", placeholder="Татарстанның башкаласы нинди шәһәр? / Какая столица Татарстана?")
+    msg = gr.Textbox(
+        label="Хәбәрегезне рус яки татар телендә языгыз",
+        placeholder="Татарстанның башкаласы нинди шәһәр? / Какая столица Татарстана?"
+    )
     clear = gr.Button("🗑️ Чистарту")
 
-    msg.submit(chat_fn, inputs=[msg, chatbot], outputs=chatbot)
-    clear.click(lambda: None, None, chatbot, queue=False)
+    msg.submit(
+        chat_fn,
+        inputs=[msg, chatbot, messages_state],
+        outputs=[chatbot, messages_state],
+    )
+
+    msg.submit(lambda: "", None, msg)
+    
+    def _reset():
+        return [], [{"role": "system", "content": SYS_PROMPT_TT}]
+
+    clear.click(_reset, inputs=None, outputs=[chatbot, messages_state], queue=False)
+    clear.click(lambda: "", None, msg, queue=False)
 
-# Запуск приложения
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))
+    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))