|
import gradio as gr |
|
from transformers import AutoTokenizer, AutoModel |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("microsoft/graphcodebert-base") |
|
model = AutoModel.from_pretrained("microsoft/graphcodebert-base") |
|
|
|
|
|
input = gr.inputs.Textbox(lines=5, label="Input") |
|
output = gr.outputs.Textbox(label="Output") |
|
|
|
|
|
def use_graphcodebert(input): |
|
|
|
input_ids = tokenizer.encode(input, return_tensors="pt") |
|
|
|
output_ids = model.generate(input_ids, max_length=50) |
|
|
|
output = tokenizer.decode(output_ids[0], skip_special_tokens=True) |
|
|
|
return output |
|
|
|
|
|
iface = gr.Interface(fn=use_graphcodebert, inputs=input, outputs=output) |
|
iface.launch() |
|
|