Spaces:
Configuration error
Configuration error
File size: 1,933 Bytes
f016346 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import gradio as gr
import torch
from transformers import RobertaTokenizer, T5ForConditionalGeneration
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
import nltk
nltk.download('punkt')
# Load model and tokenizer
model_dir = "./codet5-base-multi-sum"
tokenizer = RobertaTokenizer.from_pretrained(model_dir)
model = T5ForConditionalGeneration.from_pretrained(model_dir)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
def generate_comment(code_snippet, reference_comment):
# Add prefix for summarization task
prefixed_code = "summarize: " + code_snippet.strip()
input_ids = tokenizer(prefixed_code, return_tensors="pt").input_ids.to(device)
generated_ids = model.generate(input_ids, max_length=64, num_beams=4, early_stopping=True)
comment = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
# Tokenize and compute BLEU against user-provided reference
if reference_comment.strip():
ref_tokens = nltk.word_tokenize(reference_comment.lower())
hyp_tokens = nltk.word_tokenize(comment.lower())
bleu = sentence_bleu([ref_tokens], hyp_tokens, smoothing_function=SmoothingFunction().method1)
bleu = round(bleu, 2)
else:
bleu = "N/A (No reference provided)"
return comment, bleu
# Gradio UI
iface = gr.Interface(
fn=generate_comment,
inputs=[
gr.Textbox(label="Enter Code Snippet", lines=4, placeholder="Paste your code here..."),
gr.Textbox(label="Reference Comment (optional)", placeholder="Expected comment to compare BLEU score"),
],
outputs=[
gr.Textbox(label="Generated Comment"),
gr.Textbox(label="BLEU Score"),
],
title="Code Comment Generator using CodeT5",
description="Paste code and get a generated comment with BLEU score (optional reference)."
)
iface.launch()
|