RAVINSH commited on
Commit
db0ca24
·
verified ·
1 Parent(s): e3f5cbd

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +25 -0
app.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer
3
+
4
+ # Load the CodeGen model
5
+ model_name = "Salesforce/codegen-6B-mono"
6
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
7
+ model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
8
+
9
+ # Function to generate code
10
+ def generate_code(prompt):
11
+ inputs = tokenizer(prompt, return_tensors="pt").to("cuda") # Use GPU if available
12
+ outputs = model.generate(inputs["input_ids"], max_length=100, num_beams=5)
13
+ return tokenizer.decode(outputs[0], skip_special_tokens=True)
14
+
15
+ # Create a Gradio interface
16
+ interface = gr.Interface(
17
+ fn=generate_code,
18
+ inputs=gr.Textbox(lines=2, placeholder="Enter your code prompt here..."),
19
+ outputs="text",
20
+ title="CodeGen Code Generator",
21
+ description="Generate code from text prompts using Salesforce CodeGen-6B-mono model."
22
+ )
23
+
24
+ # Launch the app
25
+ interface.launch()