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

# =========================
# 1️⃣ Load model & tokenizer
# =========================
model_name = "Qwen/Qwen2.5-0.5B-Instruct"

# Fast tokenizer for speed
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)

# Load model with correct dtype
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",  # Uses GPU if available, else CPU
    dtype=torch.float32  # CPU inference works better with float32
)

# Optional PyTorch 2.x compile (speeds up CPU inference)
if torch.__version__.startswith("2"):
    model = torch.compile(model)

# =========================
# 2️⃣ Hardcoded system prompt
# =========================
SYSTEM_PROMPT = "You are a friendly AI assistant that gives helpful and polite answers."

# =========================
# 3️⃣ Optimized chat function
# =========================
def chat(user_prompt: str):
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_prompt}
    ]

    # Apply Qwen chat template
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    # Encode input once
    inputs = tokenizer([text], return_tensors="pt").to(model.device)

    # Faster generation settings
    outputs = model.generate(
        **inputs,
        max_new_tokens=128,   # smaller = faster
        do_sample=False,      # deterministic = faster
        num_beams=1           # no beam search
    )

    # Decode only the first sequence
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# =========================
# 4️⃣ Flask app
# =========================
app = Flask(__name__)

@app.route("/chat", methods=["GET"])
def chat_route():
    user_message = request.args.get("message")

    if not user_message:
        return jsonify({"error": "No message provided"}), 400

    try:
        response = chat(user_message)
        return jsonify({"response": response})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# =========================
# 5️⃣ Run Flask
# =========================
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)