Update app.py

app.py

@@ -25,9 +25,6 @@ with st.sidebar:
     hf_token = st.text_input("Hugging Face Token", type="password", 
                            help="Get your token from https://huggingface.co/settings/tokens")
     
-    if not hf_token:
-        st.warning("🔑 Token required for private model access!", icon="⚠️")
-    
     st.header("Upload Documents 📂")
     uploaded_file = st.file_uploader(
         "Choose a PDF or XLSX file",
@@ -42,27 +39,28 @@ if "messages" not in st.session_state:
 # Process uploaded files
 @st.cache_data
 def process_file(uploaded_file):
-    file_content = ""
+    if uploaded_file is None:
+        return ""
+    
     try:
         if uploaded_file.type == "application/pdf":
             pdf_reader = PyPDF2.PdfReader(uploaded_file)
-            file_content = "\n".join([page.extract_text() for page in pdf_reader.pages])
+            return "\n".join([page.extract_text() for page in pdf_reader.pages])
         elif uploaded_file.type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet":
             df = pd.read_excel(uploaded_file)
-            file_content = df.to_markdown()
+            return df.to_markdown()
     except Exception as e:
         st.error(f"📄 Error processing file: {str(e)}")
-    return file_content
+        return ""
 
 # Load model and tokenizer with authentication
 @st.cache_resource
 def load_model(hf_token):
     try:
-        # Login to Hugging Face Hub
         if hf_token:
             login(token=hf_token)
         else:
-            st.error("🔐 Authentication required! Please provide a Hugging Face token.")
+            st.error("🔐 Authentication required!")
             return None, None
 
         tokenizer = AutoTokenizer.from_pretrained(
@@ -82,17 +80,18 @@ def load_model(hf_token):
 
 # Generate responses with streaming
 def generate_response(prompt, file_context):
-    # Prepare prompt template
-    full_prompt = f"""You are an expert inspection engineer. Analyze this context:
+    full_prompt = f"""Analyze this context:
     {file_context}
     
     Question: {prompt}
     Answer:"""
     
-    # Create streamer
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    streamer = TextIteratorStreamer(
+        tokenizer, 
+        skip_prompt=True, 
+        skip_special_tokens=True
+    )
     
-    # Tokenize input
     inputs = tokenizer(
         full_prompt,
         return_tensors="pt",
@@ -100,7 +99,6 @@ def generate_response(prompt, file_context):
         truncation=True
     ).to(model.device)
     
-    # Start generation thread
     generation_kwargs = dict(
         inputs,
         streamer=streamer,
@@ -109,63 +107,63 @@ def generate_response(prompt, file_context):
         top_p=0.9,
         repetition_penalty=1.1,
         do_sample=True,
-        use_cache=True  # Enable caching for faster generation
+        use_cache=True
     )
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
     
+    Thread(target=model.generate, kwargs=generation_kwargs).start()
     return streamer
 
-# Display chat messages with proper avatar handling
+# Display chat messages
 for message in st.session_state.messages:
     try:
-        # Set appropriate avatar based on role
         avatar = "👤" if message["role"] == "user" else "🤖"
-        
         with st.chat_message(message["role"], avatar=avatar):
             st.markdown(message["content"])
-    except Exception as e:
-        st.error(f"Error displaying message: {str(e)}")
-        # Fallback to default avatar
+    except:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
 
-# Chat input section
+# Chat input handling
 if prompt := st.chat_input("Ask your inspection question..."):
     if not hf_token:
-        st.error("🔑 Authentication required! Please enter your Hugging Face token in the sidebar.")
+        st.error("🔑 Authentication required!")
         st.stop()
-    
+
     # Load model if not loaded
     if "model" not in st.session_state:
         st.session_state.model, st.session_state.tokenizer = load_model(hf_token)
     model = st.session_state.model
     tokenizer = st.session_state.tokenizer
     
-    # Add user message with proper avatar handling
-    try:
-        with st.chat_message("user", avatar="👤"):
-            st.markdown(prompt)
-    except:
-        # Fallback if avatar fails
-        with st.chat_message("user"):
-            st.markdown(prompt)
-    
+    # Add user message
+    with st.chat_message("user", avatar="👤"):
+        st.markdown(prompt)
     st.session_state.messages.append({"role": "user", "content": prompt})
 
     # Process file
-    file_context = process_file(uploaded_file) if uploaded_file else ""
+    file_context = process_file(uploaded_file)
     
-    # Generate and stream response with avatar handling
+    # Generate response
     if model and tokenizer:
         try:
             with st.chat_message("assistant", avatar="🤖"):
                 streamer = generate_response(prompt, file_context)
-                response = st.write_stream(streamer)
-                # Remove <think> tokens if present
-                response = response.replace("<think>", "").replace("</think>", "").strip()
-                st.session_state.messages.append({"role": "assistant", "content": response})
+                response_container = st.empty()
+                full_response = ""
+                
+                for chunk in streamer:
+                    # Remove <think> tags and clean text
+                    cleaned_chunk = chunk.replace("<think>", "").replace("</think>", "").strip()
+                    full_response += cleaned_chunk + " "
+                    
+                    # Update display with typing cursor
+                    response_container.markdown(full_response + "▌", unsafe_allow_html=True)
+                
+                # Display final response
+                response_container.markdown(full_response)
+                st.session_state.messages.append({"role": "assistant", "content": full_response})
+                
         except Exception as e:
             st.error(f"⚡ Generation error: {str(e)}")
     else:
-        st.error("🤖 Model not loaded - check your token and connection!")
+        st.error("🤖 Model not loaded!")