Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from threading import Thread
 import PyPDF2
 import pandas as pd
 import torch
@@ -8,38 +9,19 @@ import torch
 st.set_page_config(
     page_title="WizNerd Insp",
     page_icon="🚀",
-    layout="wide"
+    layout="centered"
 )
 
+# Corrected model name (fixed typo)
+MODEL_NAME = "amiguel/optimizedModelListing6.1"  # Changed from "Linsting" to "Listing"
+
 # Title with rocket emojis
 st.title("🚀 WizNerd Insp 🚀")
 
-# Define prompt template
-PROMPT_TEMPLATE = """Below is an instruction that describes a task, paired with an input that provides further context.
-You are an experienced inspection methods engineer with expertise in:
-- Offshore topside structural inspection planning
-- FLOC classification and RBI methodologies
-- Degradation mechanism analysis for process systems
-- ASME/API compliance and integrity engineering
-
-Write a response that appropriately completes the request following these steps:
-1. Analyze the context and question requirements
-2. Identify relevant codes and standards
-3. Consider equipment criticality factors
-4. Evaluate potential degradation mechanisms
-5. Formulate technical recommendation
-
-### instruction:
-{}
-
-### output:
-<think>
-{{REASONING}} 
-</think>
-{{ANSWER}}"""
-
-# Sidebar file uploader
+# Sidebar configuration
 with st.sidebar:
+    st.header("Configuration")
+    hf_token = st.text_input("HuggingFace Token", type="password")
     st.header("Upload Documents")
     uploaded_file = st.file_uploader(
         "Choose a PDF or XLSX file",
@@ -55,117 +37,90 @@ if "messages" not in st.session_state:
 @st.cache_data
 def process_file(uploaded_file):
     file_content = ""
-    
     try:
         if uploaded_file.type == "application/pdf":
             pdf_reader = PyPDF2.PdfReader(uploaded_file)
-            for page in pdf_reader.pages:
-                file_content += page.extract_text()
-                
+            file_content = "\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_string()
-            
+            file_content = df.to_markdown()
     except Exception as e:
         st.error(f"Error processing file: {str(e)}")
-        return None
-    
     return file_content
 
-# Load model and tokenizer with caching
+# Load model and tokenizer with authentication
 @st.cache_resource
 def load_model():
-    model_name = "amiguel/optimizedModelListing6.1"
-    
     try:
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_NAME,
+            token=hf_token or True
+        )
         model = AutoModelForCausalLM.from_pretrained(
-            model_name,
+            MODEL_NAME,
             device_map="auto",
             torch_dtype=torch.float16,
-            trust_remote_code=True
+            token=hf_token or True
         )
         return model, tokenizer
     except Exception as e:
-        st.error(f"Failed to load model: {str(e)}")
+        st.error(f"Model loading failed: {str(e)}")
         return None, None
 
 model, tokenizer = load_model()
 
 # Display chat messages
 for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        if message["role"] == "assistant":
-            st.markdown(message["content"]["answer"])
-            with st.expander("View Reasoning Process"):
-                st.markdown(message["content"]["reasoning"])
-        else:
-            st.markdown(message["content"])
+    with st.chat_message(message["role"], avatar="🧑💻" if message["role"] == "user" else "🤖"):
+        st.markdown(message["content"])
 
 # Chat input
 if prompt := st.chat_input("Ask your inspection question..."):
     # Add user message to chat history
+    with st.chat_message("user", avatar="🧑💻"):
+        st.markdown(prompt)
     st.session_state.messages.append({"role": "user", "content": prompt})
-    
-    # Process file if uploaded
-    file_context = ""
-    if uploaded_file is not None:
-        file_context = process_file(uploaded_file)
+
+    # Prepare context
+    file_context = process_file(uploaded_file) if uploaded_file else ""
     
     # Generate response
     if model and tokenizer:
-        with st.chat_message("assistant"):
-            with st.spinner("Analyzing..."):
-                try:
-                    # Prepare input
-                    context_prompt = f"Context: {file_context}\n\nQuestion: {prompt}" if file_context else prompt
-                    formatted_prompt = PROMPT_TEMPLATE.format(context_prompt)
-                    
-                    # Tokenize input
-                    inputs = tokenizer(
-                        formatted_prompt,
-                        return_tensors="pt",
-                        max_length=4096,
-                        truncation=True
-                    ).to(model.device)
-                    
-                    # Generate response
-                    outputs = model.generate(
-                        **inputs,
-                        max_new_tokens=1024,
-                        temperature=0.7,
-                        top_p=0.9,
-                        repetition_penalty=1.1,
-                        do_sample=True
-                    )
-                    
-                    # Decode response
-                    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-                    
-                    # Parse response components
-                    try:
-                        reasoning = full_response.split("<think>")[1].split("</think>")[0].strip()
-                        answer = full_response.split("</think>")[1].strip()
-                    except:
-                        reasoning = "Reasoning steps not properly formatted"
-                        answer = full_response
-                    
-                    # Display response
-                    with st.expander("Reasoning Process (Click to view)", expanded=False):
-                        st.markdown(f"🔍 **Analysis Steps:**\n{reasoning}")
-                    
-                    st.markdown(f"📝 **Expert Recommendation:**\n{answer}")
-                    
-                    # Add to chat history
-                    st.session_state.messages.append({
-                        "role": "assistant",
-                        "content": {
-                            "answer": answer,
-                            "reasoning": reasoning
-                        }
-                    })
-                    
-                except Exception as e:
-                    st.error(f"Generation error: {str(e)}")
+        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})
     else:
-        st.error("Model not loaded properly")
+        st.error("Model not loaded - check configuration")