|
|
import gradio as gr |
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
import torch |
|
|
|
|
|
|
|
|
model_name = "Salesforce/codegen-350M-mono" |
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
|
model = AutoModelForCausalLM.from_pretrained(model_name).to("cpu") |
|
|
|
|
|
|
|
|
def generate_code(prompt): |
|
|
inputs = tokenizer(prompt, return_tensors="pt") |
|
|
outputs = model.generate( |
|
|
**inputs, |
|
|
max_new_tokens=200, |
|
|
temperature=0.7, |
|
|
do_sample=True, |
|
|
top_k=50, |
|
|
top_p=0.95 |
|
|
) |
|
|
code = tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
|
|
|
|
|
|
|
return code[len(prompt):] if code.startswith(prompt) else code |
|
|
|
|
|
|
|
|
gr.Interface( |
|
|
fn=generate_code, |
|
|
inputs=gr.Textbox(lines=5, placeholder="Describe what code you want...", label="Prompt"), |
|
|
outputs=gr.Textbox(label="Generated Code"), |
|
|
title="Code Generator - Mono Model", |
|
|
description="Generate Python code from a text description using CodeGen-350M-Mono model" |
|
|
).launch() |
|
|
|
