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

app = Flask(__name__)

# Load fine-tuned model and tokenizer
model_path = "./finetuned_codegen"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16)

# Set padding token
tokenizer.pad_token = tokenizer.eos_token

# Move model to CPU
device = torch.device("cpu")
model.to(device)

@app.route("/", methods=["GET", "POST"])
def index():
    generated_code = ""
    prompt = ""
    if request.method == "POST":
        prompt = request.form["prompt"]
        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
        outputs = model.generate(
            **inputs,
            max_length=200,
            num_return_sequences=1,
            pad_token_id=tokenizer.eos_token_id,
            do_sample=True,
            temperature=0.7,
            top_p=0.9
        )
        generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return render_template("index.html", generated_code=generated_code, prompt=prompt)

if __name__ == "__main__":
    app.run(debug=True)