Update app.py

app.py

@@ -1,30 +1,67 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 # Загрузка токенизатора и модели
-model_name = "GoidaAlignment/GOIDA-0.5B"  # Укажите путь к вашей модели
+model_name = "GoidaAlignment/GOIDA-0.5B"  # Замените на вашу модель
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
+model = model.to("cuda" if torch.cuda.is_available() else "cpu")
 
-def generate_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
-    outputs = model.generate(inputs["input_ids"], max_length=200, num_return_sequences=1)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+# Шаблонная функция для форматирования диалога
+def apply_chat_template(chat, add_generation_prompt=True):
+    formatted_chat = ""
+    for message in chat:
+        role = message["role"]
+        content = message["content"]
+        if role == "user":
+            formatted_chat += f"User: {content}\n"
+        elif role == "assistant":
+            formatted_chat += f"Assistant: {content}\n"
+    if add_generation_prompt:
+        formatted_chat += "Assistant: "
+    return formatted_chat
 
-# Интерфейс Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("# Введите запрос, и модель ответит.")
+# Функция генерации ответа
+def generate_response(user_input, chat_history):
+    chat_history.append({"role": "user", "content": user_input})
+    formatted_chat = apply_chat_template(chat_history, add_generation_prompt=True)
+    
+    # Токенизация
+    inputs = tokenizer(formatted_chat, return_tensors="pt", add_special_tokens=False)
+    inputs = {key: tensor.to(model.device) for key, tensor in inputs.items()}
     
-    with gr.Row():
-        with gr.Column():
-            prompt_input = gr.Textbox(label="Ваш запрос", lines=4, placeholder="Введите текст")
-        with gr.Column():
-            output = gr.Textbox(label="Ответ модели", lines=6, interactive=False)
+    # Генерация
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=64,
+        temperature=0.7,
+        top_p=0.9,
+        do_sample=True
+    )
     
-    submit_button = gr.Button("Сгенерировать")
-    submit_button.click(generate_response, inputs=prompt_input, outputs=output)
+    # Декодирование
+    decoded_output = tokenizer.decode(outputs[0][inputs["input_ids"].size(1):], skip_special_tokens=True)
+    chat_history.append({"role": "assistant", "content": decoded_output})
+    
+    return decoded_output, chat_history
+
+# Интерфейс Gradio
+with gr.Blocks() as demo:
+    gr.Markdown("# Chatbot на основе модели ГОЙДАААА\nВзаимодействуйте с языковой моделью.")
+
+    chatbot = gr.Chatbot()
+    user_input = gr.Textbox(placeholder="Введите ваше сообщение...")
+    clear = gr.Button("Очистить чат")
+
+    chat_history = gr.State([])  # Состояние для хранения истории чата
+
+    user_input.submit(
+        generate_response,
+        [user_input, chat_history],
+        [chatbot, chat_history]
+    )
+    clear.click(lambda: ([], []), None, [chatbot, chat_history])
 
-# Запуск приложения
 if __name__ == "__main__":
     demo.launch()