Update app.py

app.py

@@ -4,6 +4,8 @@ from flask import Flask, render_template, request, jsonify
 from flask_cors import CORS
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from dotenv import load_dotenv
+import requests
 
 app = Flask(__name__)
 CORS(app)
@@ -39,7 +41,7 @@ def load_model():
             )
         print("Модель успешно загружена!")
     except Exception as e:
-        print("Ошибка при загрузке модели: ", e)
+        print(f"Ошибка при загрузке модели: {e}")
         model = None
         tokenizer = None
 
@@ -51,32 +53,49 @@ def before_request():
         load_model()
 
 # Функция для генерации ответа от модели
-def generate_response(prompt, max_length=1024):
-    if model is None or tokenizer is None:
-        return "Ошибка: Модель не загружена"
-    
-    # Форматирование запроса в формате Mistral Instruct
-    formatted_prompt = f"<s>[INST] {prompt} [/INST]"
-    
-    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
+from dotenv import load_dotenv
+import requests
+
+# Загружаем переменные окружения
+load_dotenv()
+
+# Проверяем наличие API ключа
+API_KEY = os.getenv("MISTRAL_API_KEY")
+API_URL = os.getenv("MISTRAL_API_URL", "https://api.mistral.ai/v1/")
+
+# Функция для генерации ответа через API
+def generate_response_api(prompt, max_length=1024):
+    if not API_KEY:
+        return "Ошибка: API ключ не найден. Пожалуйста, добавьте MISTRAL_API_KEY в файл .env"
     
-    # Генерация ответа
-    with torch.no_grad():
-        outputs = model.generate(
-            inputs["input_ids"],
-            max_new_tokens=max_length,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
+    headers = {
+        "Authorization": f"Bearer {API_KEY}",
+        "Content-Type": "application/json"
+    }
     
-    # Декодирование ответа
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Извлечение только ответа модели (после [/INST])
-    response = response.split("[/INST]")[-1].strip()
+    data = {
+        "model": "mistral-medium",  # или другая доступная модель
+        "messages": [
+            {"role": "user", "content": prompt}
+        ],
+        "max_tokens": max_length,
+        "temperature": 0.7,
+        "top_p": 0.9
+    }
     
-    return response
+    try:
+        response = requests.post(f"{API_URL}chat/completions", headers=headers, json=data)
+        response.raise_for_status()
+        result = response.json()
+        return result["choices"][0]["message"]["content"]
+    except Exception as e:
+        return f"Ошибка при обращении к API: {str(e)}"
+
+# Обновление функции generate_response для использования API
+def generate_response(prompt, max_length=1024):
+    return generate_response_api(prompt, max_length)
+
+# Удаление старой функции load_model
 
 # Маршруты
 @app.route('/')