karthik18AI commited on
Commit
e08b408
·
verified ·
1 Parent(s): 721b8a1

summary-gen

Browse files
Files changed (1) hide show
  1. app.py +20 -0
app.py ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
2
+
3
+ model_name = "Salesforce/codet5-base"
4
+ model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
5
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
6
+
7
+ def generate_summary(code_snippet):
8
+ inputs = tokenizer(code_snippet, return_tensors="pt", truncation=True, padding="max_length", max_length=512)
9
+ summary_ids = model.generate(inputs.input_ids, max_length=150, num_beams=4, length_penalty=2.0, early_stopping=True)
10
+ summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
11
+ return summary
12
+
13
+ # Example C# method
14
+ code_snippet = """
15
+ public int Add(int a, int b) {
16
+ return a + b;
17
+ }
18
+ """
19
+ summary = generate_summary(code_snippet)
20
+ print("Summary:", summary)