File size: 1,837 Bytes
5a93c7a 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 40 41 42 43 44 45 46 |
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.float32)
# 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.2, # Lower temperature for more precise outputs
top_p=0.95, # Adjusted for better sampling
top_k=50, # Added to focus on top-k tokens
no_repeat_ngram_size=3 # Prevent repetitive phrases
)
generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Clean up output to remove prompt prefix and extra text
if generated_code.startswith(prompt):
generated_code = generated_code[len(prompt):].strip()
# Remove any trailing or redundant text
generated_code = generated_code.split("\n")[0].strip() if "\n" in generated_code else generated_code
return render_template("index.html", generated_code=generated_code, prompt=prompt)
if __name__ == "__main__":
app.run(debug=True) |