import os
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"

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

app = Flask(__name__)

# 🔁 Yeni model ismi burada
model_name = "memorease/base-t5-v2"

print("[Startup] Loading model...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
print("[Startup] Model loaded.")

@app.route("/ask", methods=["POST"])
def ask_question():
    try:
        input_text = request.json.get("text")
        if not input_text:
            return jsonify({"error": "Missing 'text'"}), 400

        prompt = f"Only generate a factual and relevant question about this memory: {input_text}"
        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)

        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=64)

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

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

@app.route("/", methods=["GET"])
def healthcheck():
    return jsonify({"status": "running"})

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