Update app.py

app.py

@@ -7,7 +7,7 @@ import time
 # --- Конфигурация Hugging Face Space ---
 MODEL_NAME = "Kenan023214/PyroNet-mini"
 DEVICE = "cpu"  # Используем CPU, как указано для Basic Space
-MAX_NEW_TOKENS = 1024  # Увеличим для "хода мыслей"
+MAX_NEW_TOKENS = 512  # Увеличим для "хода мыслей"
 MAX_CONTEXT_TOKENS = 2048
 
 # Словарь с встроенным содержимым шаблонов чата
@@ -102,18 +102,14 @@ def trim_history_to_max_tokens(messages, max_tokens):
     total = 0
     kept = []
     for m in rev:
-        if len(m) == 2:
-            approx = num_tokens_of_text(m[0]) + num_tokens_of_text(m[1] if m[1] else "") + 10
-        else:
-            approx = num_tokens_of_text(m[0]) + 10
-        
+        approx = num_tokens_of_text(m["content"]) + 8
         if total + approx > max_tokens:
             break
         kept.append(m)
         total += approx
     return list(reversed(kept))
 
-def build_messages_for_template(history, reasoning: bool, language: str):
+def build_messages_for_template(history_messages, reasoning: bool, language: str):
     """Подготавливает сообщения для шаблона, включая системное сообщение."""
     full_template_content = CHAT_TEMPLATES.get(language, CHAT_TEMPLATES["en"])
     
@@ -121,12 +117,7 @@ def build_messages_for_template(history, reasoning: bool, language: str):
     system_end_tag = "<|end|>"
     system_message_raw = full_template_content.split(system_start_tag)[1].split(system_end_tag)[0].strip()
 
-    messages = [{"role": "system", "content": system_message_raw}]
-    
-    for user_msg, assistant_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        if assistant_msg is not None:
-            messages.append({"role": "assistant", "content": assistant_msg})
+    messages = [{"role": "system", "content": system_message_raw}] + list(history_messages)
 
     if reasoning:
         messages.append({"role": "user", "content": f"Режим рассуждения: покажи свои шаги, а затем окончательный ответ, начиная с '{REASONING_SEPARATOR}'"})
@@ -136,91 +127,72 @@ def build_messages_for_template(history, reasoning: bool, language: str):
 def extract_assistant_reply_and_reasoning(raw_generated_text: str) -> tuple[str, str]:
     """Убирает лишние токены и разделяет ответ на ход мыслей и окончательный ответ."""
     text = raw_generated_text
-    
-    # Remove the chat template tags
     if "<|assistant|>" in text:
         text = text.split("<|assistant|>")[-1]
     
-    # Remove end tokens
     for tag in ["<|end|>", "<|end_of_text|>", "<|end|>"]:
         text = text.replace(tag, "")
 
     text = text.strip()
 
-    # Check for the reasoning separator to prevent IndexError
     if REASONING_SEPARATOR in text:
         parts = text.split(REASONING_SEPARATOR, 1)
         reasoning = parts[0].strip()
         reply = parts[1].strip()
         return reply, reasoning
     else:
-        # If separator is not found, the whole text is the reply, and there's no reasoning.
-        return text, "Ход мыслей не предоставлен."
-
+        return text, "" # Если разделитель не найден, возвращаем все как ответ
 
 # --- Основная функция для Gradio ---
 def generate_response(user_text: str, history, reasoning: bool, language: str):
     """
     Обрабатывает пользовательский запрос, генерирует ответ и возвращает его
-    с эффектом печати и отдельным отображением хода мыслей.
+    с эффектом печати.
     """
-    try:
-        # Добавляем user-сообщение в историю Gradio
-        history.append([user_text, None])
-        
-        # 1. Первый yield: очищаем поле ввода и показываем, что модель думает
-        yield "", history, "Модель обдумывает ваш запрос..."
-
-        # Приводим историю к формату, который нужен для шаблона
-        trimmed_history = trim_history_to_max_tokens(history, MAX_CONTEXT_TOKENS)
-        messages_for_template = build_messages_for_template(trimmed_history, reasoning, language)
-        template_content = CHAT_TEMPLATES.get(language, CHAT_TEMPLATES["en"])
-        
-        # Применяем шаблон и токенизируем
-        text = tokenizer.apply_chat_template(
-            messages_for_template,
-            chat_template=template_content,
-            tokenize=False,
-            add_generation_prompt=True
+    # Добавляем user-сообщение во внутреннюю историю
+    history.append([user_text, None])
+    
+    # Конвертируем Gradio-историю в наш внутренний формат
+    internal_history = [{"role": "user", "content": h[0]} for h in history if h[0] is not None]
+    
+    trimmed_history = trim_history_to_max_tokens(internal_history, MAX_CONTEXT_TOKENS)
+    messages_for_template = build_messages_for_template(trimmed_history, reasoning, language)
+    template_content = CHAT_TEMPLATES.get(language, CHAT_TEMPLATES["en"])
+    
+    # Применяем шаблон и токенизируем
+    text = tokenizer.apply_chat_template(
+        messages_for_template,
+        chat_template=template_content,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    
+    inputs = tokenizer(text, return_tensors="pt").to(DEVICE)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=MAX_NEW_TOKENS,
+            do_sample=True,
+            top_p=0.9,
+            temperature=0.8,
+            pad_token_id=tokenizer.eos_token_id
         )
-        
-        inputs = tokenizer(text, return_tensors="pt").to(DEVICE)
 
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=MAX_NEW_TOKENS,
-                do_sample=True,
-                top_p=0.9,
-                temperature=0.8,
-                pad_token_id=tokenizer.eos_token_id
-            )
+    raw = tokenizer.decode(outputs[0], skip_special_tokens=False)
+    
+    # Извлекаем финальный ответ и ход мыслей
+    reply, reasoning_text = extract_assistant_reply_and_reasoning(raw)
 
-        raw = tokenizer.decode(outputs[0], skip_special_tokens=False)
-        
-        # Извлекаем финальный ответ и ход мыслей
-        reply, reasoning_text = extract_assistant_reply_and_reasoning(raw)
-        
-        # 2. Второй yield: обновляем поле рассуждений
-        # и начинаем "печатать" в чат. Это гарантирует, что поле рассуждений
-        # обновится сразу, до начала анимации в чате.
-        history[-1][1] = ""
+    # Обновляем историю Gradio с финальным ответом
+    history[-1][1] = ""
+    
+    # Используем генератор для создания эффекта печати
+    for chunk in reply.split():
+        history[-1][1] += chunk + " "
+        time.sleep(0.05) # Небольшая задержка для анимации
         yield "", history, reasoning_text
-
-        # 3. Последующие yield-ы для анимации печати
-        full_response = ""
-        for chunk in reply.split():
-            full_response += chunk + " "
-            history[-1][1] = full_response
-            time.sleep(0.05)
-            yield "", history, reasoning_text
-
-    except Exception as e:
-        error_message = f"Произошла ошибка: {e}"
-        history.append([None, error_message])
-        yield "", history, gr.Textbox.update(value=error_message)
-
-
+    
 # --- Интерфейс Gradio ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# PyroNet-mini Chat")