Update app.py

app.py

@@ -7,21 +7,37 @@ from sentence_transformers import SentenceTransformer
 import fitz  # PyMuPDF for better PDF extraction
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 
+
 # Configuration
 MODEL_NAME = "ibm-granite/granite-3.1-1b-a400m-instruct"
-EMBED_MODEL = "sentence-transformers/all-mpnet-base-v2"
-CHUNK_SIZE = 512
-CHUNK_OVERLAP = 64
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Initialize session state
-if "docs" not in st.session_state:
-    st.session_state.docs = []
-if "index" not in st.session_state:
-    st.session_state.index = None
-
-# Model loading with better error handling
 @st.cache_resource
+def load_model():
+    try:
+        # Load with explicit configuration
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_NAME,
+            trust_remote_code=True,
+            revision="main"
+        )
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            device_map="auto" if DEVICE == "cuda" else None,
+            torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
+            trust_remote_code=True,
+            revision="main",
+            low_cpu_mem_usage=True
+        )
+        return model, tokenizer
+    except Exception as e:
+        st.error(f"Model loading failed: {str(e)}")
+        st.stop()
+
+model, tokenizer = load_model()
+
+
 def load_models():
     try:
         tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)