Update app.py

app.py

@@ -43,52 +43,53 @@ def process_file(uploaded_file):
 
 @st.cache_resource
 def load_model(hf_token):
-    # Existing model loading logic
-    pass
+    try:
+        if not hf_token:
+            st.error("🔐 Authentication required! Please provide a Hugging Face token.")
+            return None
+        
+        # Login to Hugging Face Hub
+        login(token=hf_token)
+        
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_NAME,
+            token=hf_token
+        )
+        
+        # Load model with KV caching support
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            device_map="auto",
+            torch_dtype=torch.float16,
+            token=hf_token
+        )
+        
+        return model, tokenizer
+        
+    except Exception as e:
+        st.error(f"🤖 Model loading failed: {str(e)}")
+        return None
 
-def generate_with_kv_cache(prompt, file_context, use_cache=True):
-    full_prompt = f"Analyze this context:\n{file_context}\n\nQuestion: {prompt}\nAnswer:"
-    
-    streamer = TextIteratorStreamer(
-        tokenizer, 
-        skip_prompt=True, 
-        skip_special_tokens=True
-    )
-    
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
-    
-    # KV Caching parameters
-    generation_kwargs = {
-        **inputs,
-        "max_new_tokens": 1024,
-        "temperature": 0.7,
-        "top_p": 0.9,
-        "repetition_penalty": 1.1,
-        "do_sample": True,
-        "use_cache": use_cache,  # KV Cache control
-        "streamer": streamer
-    }
-    
-    Thread(target=model.generate, kwargs=generation_kwargs).start()
-    return streamer
-
-# Display chat messages
-for message in st.session_state.messages:
-    # Existing message display logic
-    pass
-
-# Chat input handling
+# In the main chat handling section:
 if prompt := st.chat_input("Ask your inspection question..."):
     if not hf_token:
         st.error("🔑 Authentication required!")
         st.stop()
 
-    # Load model
+    # Load model if not already loaded
     if "model" not in st.session_state:
-        st.session_state.model, st.session_state.tokenizer = load_model(hf_token)
+        model_data = load_model(hf_token)
+        if model_data is None:
+            st.error("Failed to load model. Please check your token and try again.")
+            st.stop()
+            
+        st.session_state.model, st.session_state.tokenizer = model_data
+    
     model = st.session_state.model
     tokenizer = st.session_state.tokenizer
     
+    
     # Add user message
     with st.chat_message("user", avatar="👤"):
         st.markdown(prompt)