Spaces:
Build error
Build error
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) | |