from flask import Flask, request, jsonify
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch
import os

app = Flask(__name__)

# Модель
MODEL_NAME = "KingNish/Qwen2.5-0.5b-Test-ft"

# Глобальные переменные для модели
model = None
tokenizer = None

def load_model():
    """Функция для загрузки модели при запуске приложения."""
    global model, tokenizer

    print("Loading model...")
    try:
        # Загружаем токенизатор и модель
        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        
        # Убедимся, что токенизатор имеет pad_token
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True
        )

        # Переносим модель на доступное устройство
        if torch.cuda.is_available():
            model = model.cuda()
            print("Model loaded on CUDA")
        elif hasattr(torch, 'backends') and torch.backends.mps.is_available():
            model = model.to('mps')
            print("Model loaded on MPS")
        else:
            print("Model loaded on CPU")

        print("Model loaded successfully!")

    except Exception as e:
        print(f"Error loading model: {e}")
        model = None
        tokenizer = None

# Загружаем модель при запуске приложения
load_model()

@app.route("/")
def home():
    return "API is running!"

@app.route("/health")
def health():
    return jsonify({"status": "healthy", "model_loaded": model is not None})

@app.route("/generate", methods=["POST"])
def generate():
    if model is None or tokenizer is None:
        return jsonify({"error": "Model not loaded"}), 500

    try:
        data = request.json
        prompt = data.get("prompt", "Hello")
        
        # Получаем параметры из запроса или используем значения по умолчанию
        max_length = data.get("max_length", 512)
        temperature = data.get("temperature", 0.7)
        top_p = data.get("top_p", 0.9)
        top_k = data.get("top_k", 50)
        do_sample = data.get("do_sample", True)

        # Токенизация
        inputs = tokenizer.encode(prompt, return_tensors="pt")

        # Перенос на устройство
        device = next(model.parameters()).device
        inputs = inputs.to(device)

        # Генерация
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_length=max_length,
                num_return_sequences=1,
                temperature=temperature,
                do_sample=do_sample,
                top_p=top_p,
                top_k=top_k,
                pad_token_id=tokenizer.eos_token_id,
                repetition_penalty=data.get("repetition_penalty", 1.1),
            )

        # Декодирование
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

        return jsonify({
            "prompt": prompt,
            "generated_text": generated_text,
            "status": "success"
        })

    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    app.run(host="0.0.0.0", port=port, debug=False)