| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
| tokenizer_code2desc = AutoTokenizer.from_pretrained("microsoft/codebert-base") | |
| model_code2desc = AutoModelForSeq2SeqLM.from_pretrained("microsoft/codebert-base") | |
| def code_to_description(code: str) -> str: | |
| inputs = tokenizer_code2desc.encode("summarize: " + code, return_tensors="pt", max_length=512) | |
| outputs = model_code2desc.generate(inputs, max_length=150, num_return_sequences=1, no_repeat_ngram_size=2) | |
| description = tokenizer_code2desc.decode(outputs[0], skip_special_tokens=True) | |
| return description | |
| iface = gr.Interface(fn=code_to_description, inputs="text", outputs="text") | |
| iface.launch() | |