ramky1979 commited on
Commit
0cad702
·
verified ·
1 Parent(s): 77e48b0

Create app.py

Browse files

Creating GraphCodeBERT playground.

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 AutoTokenizer, AutoModel
3
+
4
+ # Load GraphCodeBERT model and tokenizer
5
+ tokenizer = AutoTokenizer.from_pretrained("microsoft/graphcodebert-base")
6
+ model = AutoModel.from_pretrained("microsoft/graphcodebert-base")
7
+
8
+ # Define input and output interfaces
9
+ input = gr.inputs.Textbox(lines=5, label="Input")
10
+ output = gr.outputs.Textbox(label="Output")
11
+
12
+ # Define function to use GraphCodeBERT
13
+ def use_graphcodebert(input):
14
+ # Encode input
15
+ input_ids = tokenizer.encode(input, return_tensors="pt")
16
+ # Generate output
17
+ output_ids = model.generate(input_ids, max_length=50)
18
+ # Decode output
19
+ output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
20
+ # Return output
21
+ return output
22
+
23
+ # Create and launch Gradio interface
24
+ iface = gr.Interface(fn=use_graphcodebert, inputs=input, outputs=output)
25
+ iface.launch()