CodeSummarization / CodeSummarize.py
priyasaravana's picture
Upload CodeSummarize.py
52d72aa
raw
history blame
1.02 kB
import sys
from transformers import RobertaTokenizer, T5ForConditionalGeneration
def summarize_code(code_snippet):
# Load the tokenizer and model
tokenizer = RobertaTokenizer.from_pretrained('Salesforce/codet5-base')
model = T5ForConditionalGeneration.from_pretrained('Salesforce/codet5-base-multi-sum')
# Prepare the input text
input_text = code_snippet.strip()
input_ids = tokenizer.encode(input_text, return_tensors='pt')
# Generate a summary
generated_ids = model.generate(input_ids, max_length=20)
summary = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
return summary
# Example usage
# code_snippet = """
# if len(sys.argv) < 2:
# print("Usage: python <script_name.py> cluster_location image1 image2 ... imageN")
# sys.exit()
# """
code_snippet = sys.argv[1]
if len(sys.argv) < 1:
print("Usage: python <script_name.py> <code-Snippet>")
sys.exit()
summary = summarize_code(code_snippet)
print("Summary:", summary)