import flask
from flask import request, jsonify
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

app = flask.Flask(__name__)

model_id = "facebook/blenderbot-400M-distill"

print("🔄 Loading fast chat model...")

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSeq2SeqLM.from_pretrained(model_id)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

print("✅ Model loaded instantly!")

@app.route('/chat', methods=['POST'])
def chat():
    try:
        data = request.get_json()
        msg = data.get("message", "")

        if not msg:
            return jsonify({"error": "No message sent"}), 400

        inputs = tokenizer(msg, return_tensors="pt").to(device)
        output = model.generate(
            **inputs,
            max_length=200,
            do_sample=True,
            top_p=0.92,
            temperature=0.7
        )

        reply = tokenizer.decode(output[0], skip_special_tokens=True)
        return jsonify({"reply": reply})

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


if __name__ == "__main__":
    app.run(host='0.0.0.0', port=7860)