Spaces:
Build error
Build error
File size: 1,019 Bytes
52d72aa |
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 |
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)
|