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

app = Flask(__name__)

# Load the model and tokenizer
MODEL_NAME = "adityagofi/Finetunning-Gemma-2-MedicalChatbot"  # Replace with your model path
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

@app.route('/generate', methods=['POST'])
def generate():
    try:
        # Parse input data
        data = request.get_json()
        input_text = data.get("input_text", "")
        max_length = data.get("max_length", 50)
        
        # Generate text
        inputs = tokenizer.encode(input_text, return_tensors="pt")
        outputs = model.generate(inputs, max_length=max_length, num_return_sequences=1, 
                                  pad_token_id=tokenizer.eos_token_id)

        # Decode the generated text
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        return jsonify({"generated_text": generated_text})
    
    except Exception as e:
        return jsonify({"error": str(e)}), 500

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