File size: 907 Bytes
db0ca24 |
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 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
# Load the CodeGen model
model_name = "Salesforce/codegen-6B-mono"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
# Function to generate code
def generate_code(prompt):
inputs = tokenizer(prompt, return_tensors="pt").to("cuda") # Use GPU if available
outputs = model.generate(inputs["input_ids"], max_length=100, num_beams=5)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# Create a Gradio interface
interface = gr.Interface(
fn=generate_code,
inputs=gr.Textbox(lines=2, placeholder="Enter your code prompt here..."),
outputs="text",
title="CodeGen Code Generator",
description="Generate code from text prompts using Salesforce CodeGen-6B-mono model."
)
# Launch the app
interface.launch()
|