Spaces:
Sleeping
Sleeping
Update code_summarizer.py
Browse files- code_summarizer.py +2 -3
code_summarizer.py
CHANGED
|
@@ -1,10 +1,9 @@
|
|
| 1 |
from transformers import RobertaTokenizer, T5ForConditionalGeneration
|
| 2 |
|
| 3 |
-
def query(user_prompt):
|
| 4 |
tokenizer = RobertaTokenizer.from_pretrained('Salesforce/codet5-base-multi-sum')
|
| 5 |
model = T5ForConditionalGeneration.from_pretrained('Salesforce/codet5-base-multi-sum')
|
| 6 |
-
prompt=f"""<|system
|
| 7 |
-
Your task is to understand and summarize</s>\n<|user|>\n {user_prompt}\n<|assistant|>"""
|
| 8 |
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
|
| 9 |
|
| 10 |
generated_ids = model.generate(input_ids, max_length=20)
|
|
|
|
| 1 |
from transformers import RobertaTokenizer, T5ForConditionalGeneration
|
| 2 |
|
| 3 |
+
def query(lang,user_prompt):
|
| 4 |
tokenizer = RobertaTokenizer.from_pretrained('Salesforce/codet5-base-multi-sum')
|
| 5 |
model = T5ForConditionalGeneration.from_pretrained('Salesforce/codet5-base-multi-sum')
|
| 6 |
+
prompt=f"""<|system|>As an expert in {lang}, summarize the following function.</s>\n<|user|>\n {user_prompt}\n<|assistant|>"""
|
|
|
|
| 7 |
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
|
| 8 |
|
| 9 |
generated_ids = model.generate(input_ids, max_length=20)
|