Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from huggingface_hub import login
 from threading import Thread
 import PyPDF2
 import pandas as pd
@@ -12,17 +13,22 @@ st.set_page_config(
     layout="centered"
 )
 
-# Corrected model name (fixed typo)
-MODEL_NAME = "amiguel/optimizedModelListing6.1"  # Changed from "Linsting" to "Listing"
+# Correct model name
+MODEL_NAME = "amiguel/optimizedModelListing6.1"
 
 # Title with rocket emojis
 st.title("🚀 WizNerd Insp 🚀")
 
 # Sidebar configuration
 with st.sidebar:
-    st.header("Configuration")
-    hf_token = st.text_input("HuggingFace Token", type="password")
-    st.header("Upload Documents")
+    st.header("Authentication 🔒")
+    hf_token = st.text_input("HuggingFace 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",
         type=["pdf", "xlsx"],
@@ -45,29 +51,68 @@ def process_file(uploaded_file):
             df = pd.read_excel(uploaded_file)
             file_content = df.to_markdown()
     except Exception as e:
-        st.error(f"Error processing file: {str(e)}")
+        st.error(f"📄 Error processing file: {str(e)}")
     return file_content
 
 # Load model and tokenizer with authentication
 @st.cache_resource
-def load_model():
+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.")
+            return None, None
+
         tokenizer = AutoTokenizer.from_pretrained(
             MODEL_NAME,
-            token=hf_token or True
+            token=hf_token
         )
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
             device_map="auto",
             torch_dtype=torch.float16,
-            token=hf_token or True
+            token=hf_token
         )
         return model, tokenizer
     except Exception as e:
-        st.error(f"Model loading failed: {str(e)}")
+        st.error(f"🤖 Model loading failed: {str(e)}")
         return None, None
 
-model, tokenizer = load_model()
+# 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:
+    {file_context}
+    
+    Question: {prompt}
+    Answer:"""
+    
+    # Create streamer
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
+    
+    # Tokenize input
+    inputs = tokenizer(
+        full_prompt,
+        return_tensors="pt",
+        max_length=4096,
+        truncation=True
+    ).to(model.device)
+    
+    # Start generation thread
+    generation_kwargs = dict(
+        inputs,
+        streamer=streamer,
+        max_new_tokens=1024,
+        temperature=0.7,
+        top_p=0.9,
+        repetition_penalty=1.1
+    )
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    
+    return streamer
 
 # Display chat messages
 for message in st.session_state.messages:
@@ -76,51 +121,32 @@ for message in st.session_state.messages:
 
 # Chat input
 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.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 to chat history
     with st.chat_message("user", avatar="🧑💻"):
         st.markdown(prompt)
     st.session_state.messages.append({"role": "user", "content": prompt})
 
-    # Prepare context
+    # Process file
     file_context = process_file(uploaded_file) if uploaded_file else ""
     
-    # Generate response
+    # Generate and stream response
     if model and tokenizer:
         with st.chat_message("assistant", avatar="🤖"):
-            # Prepare prompt template
-            full_prompt = f"""You are an expert inspection engineer. Analyze this context:
-            {file_context}
-            
-            Question: {prompt}
-            Answer:"""
-            
-            # Create streamer
-            streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
-            
-            # Tokenize input
-            inputs = tokenizer(
-                full_prompt,
-                return_tensors="pt",
-                max_length=4096,
-                truncation=True
-            ).to(model.device)
-            
-            # Start generation thread
-            generation_kwargs = dict(
-                inputs,
-                streamer=streamer,
-                max_new_tokens=1024,
-                temperature=0.7,
-                top_p=0.9,
-                repetition_penalty=1.1
-            )
-            thread = Thread(target=model.generate, kwargs=generation_kwargs)
-            thread.start()
-            
-            # Stream response
-            response = st.write_stream(streamer)
-            
-            # Add to chat history
-            st.session_state.messages.append({"role": "assistant", "content": response})
+            try:
+                streamer = generate_response(prompt, file_context)
+                response = st.write_stream(streamer)
+                st.session_state.messages.append({"role": "assistant", "content": response})
+            except Exception as e:
+                st.error(f"⚡ Generation error: {str(e)}")
     else:
-        st.error("Model not loaded - check configuration")
+        st.error("🤖 Model not loaded - check your token and connection!")