Update app.py

app.py

@@ -1,10 +1,9 @@
 import gradio as gr
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import random
-from flask import Flask, request, jsonify
 
 # Загрузка модели и токенизатора
-model_name = "ai-forever/rugpt3small_based_on_gpt2"
+model_name = "ai-forever/rugpt3large_based_on_gpt2"
 tokenizer = GPT2Tokenizer.from_pretrained(model_name)
 model = GPT2LMHeadModel.from_pretrained(model_name)
 
@@ -14,14 +13,14 @@ with open("dialogues.txt", "r", encoding="utf-8") as file:
 
 # Функция генерации ответа
 def generate(prompt, _=None):
-    user_input = f"[USER]: {prompt}"
+    user_input = f"[USER]: {prompt}\n[BOT]: {random.choice(random_phrases)},"
     inputs = tokenizer.encode(prompt, return_tensors="pt")
     outputs = model.generate(inputs, max_length=50, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
     # Прекращаем генерацию на строке [USER]
-    if "[USER]" in response:
-        response = response.split("[USER]")[0]
+    if "[" in response:
+        response = response.split("[")[0]
     
     bot_message = f"[BOT]: {random.choice(random_phrases)}, {response.strip()}"
     return bot_message
@@ -40,20 +39,5 @@ demo = gr.ChatInterface(
     undo_btn=None
 )
 
-# Flask приложение
-app = Flask(__name__)
-
-@app.route("/generate", methods=["POST"])
-def generate_api():
-    data = request.json
-    if "prompt" not in data:
-        return jsonify({"error": "На что я должен отвечать, гений?"}), 400
-    
-    prompt = data["prompt"]
-    response = generate(prompt)
-    return jsonify({"response": response})
-
-# Запуск Gradio и Flask приложений
 if __name__ == "__main__":
     demo.launch(share=True)
-    app.run(host="0.0.0.0", port=5000)