Update app.py

app.py

@@ -1,119 +1,119 @@
-import os
-import json
-from flask import Flask, render_template, request, jsonify
-from flask_cors import CORS
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-app = Flask(__name__)
-CORS(app)
-
-# Глобальные переменные для модели
-model = None
-tokenizer = None
-
-# Загрузка модели Mistral (локально)
-def load_model():
-    global model, tokenizer
-    try:
-        model_name = "mistralai/Mistral-7B-Instruct-v0.3"
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        
-        # Проверяем наличие GPU
-        if torch.cuda.is_available():
-            print("Загрузка модели на GPU...")
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                torch_dtype=torch.float16,
-                device_map="auto",
-                load_in_8bit=True  # Для оптимизации памяти
-            )
-        else:
-            print("GPU не обнаружен. Загрузка модели на CPU (это может быть медленно)...")
-            # Загрузка облегченной версии для CPU
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                torch_dtype=torch.float32,  # Используем float32 для CPU
-                low_cpu_mem_usage=True,
-                device_map="auto"
-            )
-        print("Модель успешно загружена!")
-    except Exception as e:
-        print(f"Ошибка при загрузке модели: {e}")
-        model = None
-        tokenizer = None
-
-# Инициализация приложения
-@app.before_request
-def before_request():
-    global model, tokenizer
-    if model is None or tokenizer is None:
-        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)
-    
-    # Генерация ответа
-    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
-        )
-    
-    # Декодирование ответа
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Извлечение только ответа модели (после [/INST])
-    response = response.split("[/INST]")[-1].strip()
-    
-    return response
-
-# Маршруты
-@app.route('/')
-def index():
-    return render_template('index.html')
-
-@app.route('/api/chat', methods=['POST'])
-def chat():
-    data = request.json
-    prompt = data.get('prompt', '')
-    
-    if not prompt:
-        return jsonify({"error": "Пустой запрос"}), 400
-    
-    try:
-        response = generate_response(prompt)
-        return jsonify({"response": response})
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
-
-@app.route('/api/code', methods=['POST'])
-def code():
-    data = request.json
-    prompt = data.get('prompt', '')
-    language = data.get('language', 'python')
-    
-    if not prompt:
-        return jsonify({"error": "Пустой запрос"}), 400
-    
-    # Добавляем контекст для генерации кода
-    code_prompt = f"Напиши код на языке {language} для решения следующей задачи: {prompt}"
-    
-    try:
-        response = generate_response(code_prompt)
-        return jsonify({"code": response})
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
-
-if __name__ == '__main__':
+import os
+import json
+from flask import Flask, render_template, request, jsonify
+from flask_cors import CORS
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+app = Flask(__name__)
+CORS(app)
+
+# Глобальные переменные для модели
+model = None
+tokenizer = None
+
+# Загрузка модели Mistral (локально)
+def load_model():
+    global model, tokenizer
+    try:
+        model_name = "mistralai/Mistral-7B-Instruct-v0.3"
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+        # Проверяем наличие GPU
+        if torch.cuda.is_available():
+            print("Загрузка модели на GPU...")
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                load_in_8bit=True  # Для оптимизации памяти
+            )
+        else:
+            print("GPU не обнаружен. Загрузка модели на CPU (это может быть медленно)...")
+            # Загрузка облегченной версии для CPU
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float32,  # Используем float32 для CPU
+                low_cpu_mem_usage=True,
+                device_map="auto"
+            )
+        print("Модель успешно загружена!")
+    except Exception as e:
+        print("Ошибка при загрузке модели: ", e)
+        model = None
+        tokenizer = None
+
+# Инициализация приложения
+@app.before_request
+def before_request():
+    global model, tokenizer
+    if model is None or tokenizer is None:
+        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)
+    
+    # Генерация ответа
+    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
+        )
+    
+    # Декодирование ответа
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Извлечение только ответа модели (после [/INST])
+    response = response.split("[/INST]")[-1].strip()
+    
+    return response
+
+# Маршруты
+@app.route('/')
+def index():
+    return render_template('index.html')
+
+@app.route('/api/chat', methods=['POST'])
+def chat():
+    data = request.json
+    prompt = data.get('prompt', '')
+    
+    if not prompt:
+        return jsonify({"error": "Пустой запрос"}), 400
+    
+    try:
+        response = generate_response(prompt)
+        return jsonify({"response": response})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+@app.route('/api/code', methods=['POST'])
+def code():
+    data = request.json
+    prompt = data.get('prompt', '')
+    language = data.get('language', 'python')
+    
+    if not prompt:
+        return jsonify({"error": "Пустой запрос"}), 400
+    
+    # Добавляем контекст для генерации кода
+    code_prompt = f"Напиши код на языке {language} для решения следующей задачи: {prompt}"
+    
+    try:
+        response = generate_response(code_prompt)
+        return jsonify({"code": response})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+if __name__ == '__main__':
     app.run(debug=True, port=5000)