Update app.py

app.py

@@ -4,7 +4,7 @@ st.set_page_config(page_title="RAG Book Analyzer", layout="wide")  # Must be the
 import torch
 import numpy as np
 import faiss
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer
 import fitz  # PyMuPDF for PDF extraction
 import docx2txt  # For DOCX extraction
@@ -13,8 +13,8 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
 # ------------------------
 # Configuration
 # ------------------------
-MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
-EMBED_MODEL = "sentence-transformers/all-mpnet-base-v2"
+MODEL_NAME = "microsoft/phi-2"  # Open-source model with good performance
+EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"  # Smaller embedding model
 CHUNK_SIZE = 512
 CHUNK_OVERLAP = 64
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -25,12 +25,12 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 @st.cache_resource
 def load_models():
     try:
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
             device_map="auto" if DEVICE == "cuda" else None,
             torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
-            low_cpu_mem_usage=True
+            trust_remote_code=True
         )
         embedder = SentenceTransformer(EMBED_MODEL, device=DEVICE)
         return tokenizer, model, embedder
@@ -83,8 +83,8 @@ def build_index(chunks):
 # Summarization and Q&A Functions
 # ------------------------
 def generate_summary(text):
-    # Create prompt with Mistral format
-    prompt = f"<s>[INST] Summarize this book in a concise paragraph: {text[:3000]} [/INST]"
+    # Create prompt for Phi-2 model
+    prompt = f"Instruct: Summarize this book in a concise paragraph\nInput: {text[:3000]}\nOutput:"
     inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
     outputs = model.generate(
         **inputs,
@@ -94,11 +94,11 @@ def generate_summary(text):
         do_sample=True
     )
     summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return summary.split("[/INST]")[-1].strip()
+    return summary.split("Output:")[-1].strip()
 
 def generate_answer(query, context):
-    # Create prompt with Mistral format
-    prompt = f"<s>[INST] Answer this question based on the context. If unsure, say 'I don't know'.\n\nQuestion: {query}\nContext: {context} [/INST]"
+    # Create prompt for Phi-2 model
+    prompt = f"Instruct: Answer this question based on the context. If unsure, say 'I don't know'.\nQuestion: {query}\nContext: {context}\nOutput:"
     inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
     outputs = model.generate(
         **inputs,
@@ -109,7 +109,7 @@ def generate_answer(query, context):
         do_sample=True
     )
     answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return answer.split("[/INST]")[-1].strip()
+    return answer.split("Output:")[-1].strip()
 
 # ------------------------
 # Streamlit UI