Update app.py

app.py

@@ -1,17 +1,12 @@
-import os
-from flask import Flask, request, jsonify
+import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from huggingface_hub import login
 
-app = Flask(__name__)
-
 def init_model():
     global model, tokenizer
-    hf_token = os.getenv("HF_TOKEN")  # Чтение токена из переменной окружения
-
-    if hf_token is None:
-        raise ValueError("Hugging Face token is not set. Please set the HF_TOKEN environment variable.")
+    # Вставьте сюда ваш токен доступа Hugging Face
+    hf_token = os.getenv("HF_TOKEN")
 
     # Аутентификация с использованием токена
     login(hf_token, add_to_git_credential=True)
@@ -29,18 +24,8 @@ def init_model():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
 
-@app.route('/generate', methods=['POST'])
-def generate_response():
+def generate_response(prompt, max_length=100, temperature=0.7, top_p=0.85, repetition_penalty=1.1):
     try:
-        data = request.get_json()
-        print(f"Received data: {data}")
-
-        prompt = data.get('prompt', '')
-        max_length = data.get('max_length', 100)
-        temperature = data.get('temperature', 0.7)
-        top_p = data.get('top_p', 0.85)
-        repetition_penalty = data.get('repetition_penalty', 1.1)
-
         input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
         attention_mask = torch.ones_like(input_ids).to(model.device)
 
@@ -55,15 +40,28 @@ def generate_response():
             num_return_sequences=1,
             pad_token_id=tokenizer.eos_token_id
         )
-        print(f"Generated output: {output}")
         response_text = tokenizer.decode(output[0], skip_special_tokens=True)
-        print(f"Generated response: {response_text}")
-        return jsonify({"response": response_text})
-
+        return response_text
     except Exception as e:
-        print(f"Error: {str(e)}")
-        return jsonify({"response": "Извините, произошла ошибка при генерации ответа."}), 500
+        return f"Извините, произошла ошибка при генерации ответа: {str(e)}"
+
+# Инициализация модели и токенизатора
+init_model()
+
+# Создание интерфейса Gradio
+iface = gr.Interface(
+    fn=generate_response,
+    inputs=[
+        gr.inputs.Textbox(lines=2, placeholder="Введите ваш текст здесь..."),
+        gr.inputs.Slider(20, 200, step=1, default=100, label="Максимальная длина"),
+        gr.inputs.Slider(0.1, 1.0, step=0.1, default=0.7, label="Температура"),
+        gr.inputs.Slider(0.1, 1.0, step=0.05, default=0.85, label="Top-p"),
+        gr.inputs.Slider(1.0, 2.0, step=0.1, default=1.1, label="Штраф за повторение")
+    ],
+    outputs="text",
+    title="LLM Model Demo",
+    description="Введите текстовый запрос, чтобы сгенерировать ответ с помощью LLM модели."
+)
 
 if __name__ == "__main__":
-    init_model()
-    app.run(host='0.0.0.0', port=7860)
+    iface.launch()