Update app.py

app.py

@@ -8,46 +8,71 @@ print("Загружаем модель и токенизатор...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
-# Если в Spaces доступен GPU, переводим модель на него
-device = 0 if torch.cuda.is_available() else -1
-if device == 0:
-    model = model.to("cuda")
+# Убедимся, что токенизатор и модель используют одинаковый словарь
+tokenizer.pad_token = tokenizer.eos_token
+model.config.pad_token_id = model.config.eos_token_id
 
+# Если в Spaces доступен GPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device)
 model.eval()
 
 def chat(user_input):
-    # Дополнительно можно добавить обозначение конца строки для корректного завершения генерации
-    input_with_eos = user_input + tokenizer.eos_token
+    # Формируем промпт
+    prompt = f"User: {user_input}\nAssistant:"
     
-    # Токенизируем входной текст
-    inputs = tokenizer.encode(input_with_eos, return_tensors="pt")
-    if device >= 0:
-        inputs = inputs.to("cuda")
+    try:
+        # Токенизируем с явным указанием параметров
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=512,  # Ограничиваем длину входного текста
+            add_special_tokens=True
+        )
         
-    # Генерация ответа
-    outputs = model.generate(
-        inputs,
-        max_length=200,                 # можно изменить длину генерируемого текста
-        pad_token_id=tokenizer.eos_token_id,
-        do_sample=True,
-        top_p=0.9,
-        temperature=0.7
-    )
-    # Декодируем сгенерированный текст
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Если модель возвращает и исходный текст, можно убрать его:
-    if generated_text.startswith(user_input):
-        generated_text = generated_text[len(user_input):].strip()
-    return generated_text
+        # Переносим тензоры на нужное устройство
+        input_ids = inputs["input_ids"].to(device)
+        attention_mask = inputs["attention_mask"].to(device)
+
+        # Генерация с обработкой ошибок
+        with torch.no_grad():
+            outputs = model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                max_length=200,  # Ограничиваем длину выходного текста
+                pad_token_id=tokenizer.eos_token_id,
+                do_sample=True,
+                top_p=0.9,
+                temperature=0.7,
+                num_return_sequences=1,
+                no_repeat_ngram_size=3  # Избегаем повторений
+            )
+
+        # Декодируем результат
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Убираем исходный промпт из ответа
+        response = generated_text.split("Assistant:")[-1].strip()
+        
+        return response if response else "Извините, не удалось сгенерировать ответ."
+        
+    except Exception as e:
+        print(f"Ошибка при генерации: {str(e)}")
+        return f"Произошла ошибка при обработке запроса: {str(e)}"
 
 # Создаем интерфейс Gradio
 iface = gr.Interface(
     fn=chat,
-    inputs=gr.Textbox(lines=2, placeholder="Например: Привет, как дела?", label="Введите сообщение"),
+    inputs=gr.Textbox(
+        lines=2,
+        placeholder="Например: Привет, как дела?",
+        label="Введите сообщение"
+    ),
     outputs=gr.Textbox(label="Ответ модели"),
     title="RuDialoGPT-small Chat",
-    description="Диалоговый чат на базе модели t-bank-ai/RuDialoGPT-small от Hugging Face"
+    description="Диалоговый чат на базе модели t-bank-ai/RuDialoGPT-small"
 )
 
 if __name__ == "__main__":