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

app = Flask(__name__)

# Load model and tokenizer once at startup
model_name = "gpt-oss-20b"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

@app.route("/generate", methods=["GET"])
def generate_text():
    # Get 'message' parameter from query string
    prompt = request.args.get("message", "")
    if not prompt:
        return jsonify({"error": "Please provide a 'message' parameter."}), 400

    # Tokenize input
    inputs = tokenizer(prompt, return_tensors="pt")

    # Generate text
    outputs = model.generate(
        **inputs,
        max_length=100,
        num_return_sequences=1,
        no_repeat_ngram_size=2,
        temperature=0.7
    )

    # Decode generated text
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Return as JSON
    return jsonify({"input": prompt, "output": generated_text})

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