Update qa_summary.py

qa_summary.py

@@ -3,7 +3,7 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 @spaces.GPU(duration=60)
-def generate_answer(llm_name, texts, query, queries, mode='validate'):
+def generate_answer(llm_name, texts, query, queries, response_lang, mode='validate'):
     
     if llm_name == 'solar':
         tokenizer = AutoTokenizer.from_pretrained("Upstage/SOLAR-10.7B-Instruct-v1.0", use_fast=True)
@@ -17,7 +17,7 @@ def generate_answer(llm_name, texts, query, queries, mode='validate'):
         tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2", use_fast=True)
         llm_model = AutoModelForCausalLM.from_pretrained(
         "mistralai/Mistral-7B-Instruct-v0.2",
-        #device_map="auto", 
+        # device_map="auto", 
         device_map="cuda",
         torch_dtype=torch.float16,
         )
@@ -28,7 +28,7 @@ def generate_answer(llm_name, texts, query, queries, mode='validate'):
         "microsoft/Phi-3-mini-128k-instruct", 
         device_map="auto", 
         torch_dtype="auto", 
-        trust_remote_code=False,
+        trust_remote_code=True,
         )  
     
     template_texts =""
@@ -42,8 +42,8 @@ def generate_answer(llm_name, texts, query, queries, mode='validate'):
     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. Also, return the top 5 unaltered documents that answer the queries.\nQueries: {queries} \nDocuments: {template_texts}.'} ]
-        conversation = [ {'role': 'user', 'content': f'For the following queries and documents, in a brief paragraph try to answer the given queries based on the documents. Then, return the top 5 documents as provided that answer the queries.\nQueries: {queries} \nDocuments: {template_texts}.'} ]
+        conversation = [ {'role': 'user', 'content': f"""For the following queries and documents, in a brief paragraph try to answer the given queries based on the documents. 
+        Then, return the top 5 documents as provided that answer the queries.\nQueries: {queries} \nDocuments: {template_texts}. Give your response in {response_lang} language"""} ]
         
 
     prompt = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)