|
|
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_length=128, num_return_sequences=1) |
|
|
code = tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
|
return 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() |
|
|
|
