Update app.py

app.py

@@ -20,6 +20,24 @@ import torch
 import tqdm 
 import accelerate
 
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+translation_model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+translation_tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+
+def english_to_hindi(sentence):
+    translation_tokenizer.src_lang = "en_xx"
+    encoded_hi = translation_tokenizer(sentence, return_tensors="pt")
+    generated_tokens = translation_model.generate(**encoded_hi, forced_bos_token_id=translation_tokenizer.lang_code_to_id["hi_IN"] )
+    return (translation_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True))
+
+
+def hindi_to_english(sentence):
+    translation_tokenizer.src_lang = "hi_IN"
+    encoded_hi = translation_tokenizer(sentence, return_tensors="pt")
+    generated_tokens = translation_model.generate(**encoded_hi, forced_bos_token_id=translation_tokenizer.lang_code_to_id["en_XX"] )
+    return (translation_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True))
+
+
 llm_model = "mistralai/Mistral-7B-Instruct-v0.2"
 
 
@@ -162,7 +180,8 @@ def conversation(qa_chain, message, history):
    
     # Generate response using QA chain
     response = qa_chain({"question": message, "chat_history": formatted_chat_history})
-    response_answer = response["answer"]
+    #response_answer = response["answer"]
+    response_answer = english_to_hindi(response["answer"])[0]
     if response_answer.find("Helpful Answer:") != -1:
         response_answer = response_answer.split("Helpful Answer:")[-1]
     response_sources = response["source_documents"]