dialo

app.py

@@ -1,29 +1,50 @@
 import os
 import gradio as gr
-from transformers import GPT2Tokenizer, GPT2LMHeadModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 # Импортируем токены из переменных окружения
 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
 
-# Загрузка модели и токенизатора с использованием токена
-model_name = "gpt2"  # Или другой, если нужно
-tokenizer = GPT2Tokenizer.from_pretrained(model_name, use_auth_token=HUGGINGFACE_TOKEN)
-model = GPT2LMHeadModel.from_pretrained(model_name, use_auth_token=HUGGINGFACE_TOKEN)
+# Загрузка модели и токенизатора DialoGPT
+model_name = "microsoft/DialoGPT-medium"  # Можно использовать small или large версии
+tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_TOKEN)
+model = AutoModelForCausalLM.from_pretrained(model_name, token=HUGGINGFACE_TOKEN)
 
 # Функция для ведения диалога
-def chat_with_model(user_input):
-    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
-    chat_history_ids = model.generate(input_ids, max_length=100, pad_token_id=tokenizer.eos_token_id)
-    bot_response = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
-    return bot_response
+def chat_with_model(user_input, chat_history=[]):
+    # Кодируем входное сообщение и добавляем к истории
+    new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
+
+    # Объединяем новую информацию с историей
+    bot_input_ids = torch.cat([torch.tensor(chat_history), new_user_input_ids], dim=-1) if chat_history else new_user_input_ids
+
+    # Генерируем ответ
+    chat_history_ids = model.generate(
+        bot_input_ids,
+        max_length=100,
+        num_return_sequences=1,
+        pad_token_id=tokenizer.eos_token_id,
+        temperature=0.7,
+        top_k=50,
+        top_p=0.95
+    )
+
+    # Декодируем ответ
+    bot_response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
+
+    # Обновляем историю
+    chat_history.append(new_user_input_ids)
+    chat_history.append(chat_history_ids[:, bot_input_ids.shape[-1]:])
+
+    return bot_response, chat_history
 
 # Создание интерфейса Gradio
 iface = gr.Interface(
     fn=chat_with_model,
     inputs="text",
     outputs="text",
-    title="Чатбот на GPT-2",
+    title="Чатбот на DialoGPT",
     description="Поболтай со своим чатботом!"
 )