File size: 1,285 Bytes
6e793bd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
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) |