Update app.py

app.py

@@ -12,18 +12,16 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.docstore.document import Document
 import PyPDF2
 import tempfile
-from bitsandbytes import functional as F
-import bitsandbytes as bnb
 
 # Initialize BLIP for image captioning
 blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
 blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large").to("cuda" if torch.cuda.is_available() else "cpu")
 
-# Initialize Mistral-7B for conversational tasks
-mistral_model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-mistral_tokenizer = AutoTokenizer.from_pretrained(mistral_model_name)
-mistral_model = AutoModelForCausalLM.from_pretrained(
-    mistral_model_name,
+# Initialize LLaMA-3-8B-Instruct for conversational tasks
+llama_model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
+llama_tokenizer = AutoTokenizer.from_pretrained(llama_model_name)
+llama_model = AutoModelForCausalLM.from_pretrained(
+    llama_model_name,
     load_in_4bit=True,
     device_map="auto",
     torch_dtype=torch.bfloat16
@@ -104,11 +102,11 @@ def chat_with_llm(message, history):
     if vector_store:
         docs = vector_store.similarity_search(message, k=3)
         context = "\n".join([doc.page_content for doc in docs])
-    prompt = f"[INST] You are a helpful assistant. Use the following context to answer the question accurately:\n\n{context}\n\nQuestion: {message} [/INST]"
-    inputs = mistral_tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
+    prompt = f"<|begin_of_text|><|start_header_id|>system<|end_header_id>You are a helpful assistant. Use the following context to answer the question accurately:\n\n{context}<|eot_id><|start_header_id|>user<|end_header_id>{message}<|eot_id><|start_header_id|>assistant<|end_header_id>"
+    inputs = llama_tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
     with torch.no_grad():
-        outputs = mistral_model.generate(**inputs, max_length=1000, num_return_sequences=1, temperature=0.7)
-    response = mistral_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        outputs = llama_model.generate(**inputs, max_length=1000, num_return_sequences=1, temperature=0.7)
+    response = llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response.replace(prompt, "").strip()
 
 def image_tab(image, target_languages):