| 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() |
|
|