sombochea commited on
Commit
8ab6f65
1 Parent(s): cae8e67

Add application file

Browse files
Files changed (1) hide show
  1. app.py +39 -0
app.py ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import AutoModelForCausalLM, AutoTokenizer
4
+
5
+ # Load model and tokenizer
6
+ tokenizer = AutoTokenizer.from_pretrained(
7
+ "stabilityai/stable-code-3b", trust_remote_code=True)
8
+ model = AutoModelForCausalLM.from_pretrained(
9
+ "stabilityai/stable-code-3b",
10
+ trust_remote_code=True,
11
+ torch_dtype="auto",
12
+ ).to("cuda" if torch.cuda.is_available() else "cpu") # Check for GPU availability
13
+
14
+ # Define the Gradio interface
15
+ iface = gr.Interface(
16
+ fn=generate_code,
17
+ inputs=[gr.Textbox(lines=2, placeholder="Enter your Python code prompt")],
18
+ outputs="textbox",
19
+ title="Python Code Completion",
20
+ description="Generate code completions using a large language model.",
21
+ )
22
+
23
+ # Define the main function for code generation
24
+
25
+
26
+ def generate_code(prompt):
27
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
28
+ tokens = model.generate(
29
+ **inputs,
30
+ max_new_tokens=48,
31
+ temperature=0.2,
32
+ do_sample=True,
33
+ )
34
+ generated_code = tokenizer.decode(tokens[0], skip_special_tokens=True)
35
+ return generated_code
36
+
37
+
38
+ # Launch the Gradio app
39
+ iface.launch()