app.py

import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
from threading import Thread
import PyPDF2
import pandas as pd
import torch

# Set page configuration
st.set_page_config(
    page_title="WizNerd Insp",
    page_icon="🚀",
    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 🚀")

# 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",
        type=["pdf", "xlsx"],
        label_visibility="collapsed"
    )

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Process uploaded files
@st.cache_data
def process_file(uploaded_file):
    file_content = ""
    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])
        elif uploaded_file.type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet":
            df = pd.read_excel(uploaded_file)
            file_content = df.to_markdown()
    except Exception as e:
        st.error(f"Error processing file: {str(e)}")
    return file_content

# Load model and tokenizer with authentication
@st.cache_resource
def load_model():
    try:
        tokenizer = AutoTokenizer.from_pretrained(
            MODEL_NAME,
            token=hf_token or True
        )
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            device_map="auto",
            torch_dtype=torch.float16,
            token=hf_token or True
        )
        return model, tokenizer
    except Exception as 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"], 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})

    # Prepare context
    file_context = process_file(uploaded_file) if uploaded_file else ""
    
    # Generate 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})
    else:
        st.error("Model not loaded - check configuration")