priyasaravana commited on
Commit
52d72aa
1 Parent(s): 6985dfb

Upload CodeSummarize.py

Browse files
Files changed (1) hide show
  1. CodeSummarize.py +34 -0
CodeSummarize.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import sys
2
+
3
+ from transformers import RobertaTokenizer, T5ForConditionalGeneration
4
+
5
+ def summarize_code(code_snippet):
6
+ # Load the tokenizer and model
7
+ tokenizer = RobertaTokenizer.from_pretrained('Salesforce/codet5-base')
8
+ model = T5ForConditionalGeneration.from_pretrained('Salesforce/codet5-base-multi-sum')
9
+
10
+ # Prepare the input text
11
+ input_text = code_snippet.strip()
12
+ input_ids = tokenizer.encode(input_text, return_tensors='pt')
13
+
14
+ # Generate a summary
15
+ generated_ids = model.generate(input_ids, max_length=20)
16
+ summary = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
17
+
18
+ return summary
19
+
20
+ # Example usage
21
+ # code_snippet = """
22
+
23
+ # if len(sys.argv) < 2:
24
+ # print("Usage: python <script_name.py> cluster_location image1 image2 ... imageN")
25
+ # sys.exit()
26
+
27
+ # """
28
+ code_snippet = sys.argv[1]
29
+ if len(sys.argv) < 1:
30
+ print("Usage: python <script_name.py> <code-Snippet>")
31
+ sys.exit()
32
+
33
+ summary = summarize_code(code_snippet)
34
+ print("Summary:", summary)