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)