Update qa_summary.py

qa_summary.py

@@ -1,7 +1,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
-def generate_answer(llm_name, texts, query, mode='validate'):
+def generate_answer(llm_name, texts, query, queries, mode='validate'):
     
     if llm_name == 'solar':
         tokenizer = AutoTokenizer.from_pretrained("Upstage/SOLAR-10.7B-Instruct-v1.0", use_fast=True)
@@ -38,6 +38,9 @@ def generate_answer(llm_name, texts, query, mode='validate'):
         conversation = [ {'role': 'user', 'content': f'For the following query and documents, try to answer the given query based on the documents.\nQuery: {query} \nDocuments: {template_texts}.'} ]
     elif mode == 'h_summarize':
         conversation = [ {'role': 'user', 'content': f'The documents below describe a developing disaster event. Based on these documents, write a brief summary in the form of a paragraph, highlighting the most crucial information. \nDocuments: {template_texts}'} ]
+    elif mode == "multi_summarize":
+        conversation = [ {'role': 'user', 'content': f'For the following queries and documents, try to answer the given queries based on the documents.\nQueries: {queries} \nDocuments: {template_texts}.'} ]
+        
 
     prompt = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer(prompt, return_tensors="pt").to(llm_model.device)