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