Update app.py

app.py

@@ -1,612 +1,156 @@
+# ✅ SmartManuals-AI app.py (for Hugging Face Spaces)
+# Optimized to support multiple LLMs, Gradio UI, and secure on-device document QA
+
 import os
 import json
-import fitz  # PyMuPDF
+import io
+import fitz
 import nltk
 import chromadb
+import pytesseract
+import numpy as np
+import torch
+from PIL import Image
 from tqdm import tqdm
 from nltk.tokenize import sent_tokenize
 from sentence_transformers import SentenceTransformer, util
-import numpy as np
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import pytesseract
-from PIL import Image
-import io
 import gradio as gr
 
-# ---------------------------
-# ⚙️ Configuration
-# ---------------------------
-pdf_folder = r"./Manuals" # Path relative to the app.py file in the Space
-output_jsonl_pages = "manual_pages_with_ocr.jsonl"
-output_jsonl_chunks = "manual_chunks_with_ocr.jsonl"
-chroma_path = "./chroma_store"
-collection_name = "manual_chunks"
-chunk_size = 750
-chunk_overlap = 100
-MAX_CONTEXT_CHUNKS = 3 # Max chunks to send to the LLM
-
-# Hugging Face Model Configuration
-HF_MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
-# Read HF Token from environment variable for security
-HF_TOKEN = os.environ.get("HF_TOKEN") # Hugging Face Space secret name
-
-# ---------------------------
-# Ensure NLTK resources are available
-# ---------------------------
+# ----------------------
+# 🔧 Configurations
+# ----------------------
+PDF_DIR = "./Manuals"
+CHROMA_PATH = "./chroma_store"
+COLLECTION_NAME = "manual_chunks"
+MAX_CONTEXT_CHUNKS = 3
+CHUNK_SIZE = 750
+CHUNK_OVERLAP = 100
+MODEL_OPTIONS = [
+    "meta-llama/Llama-3.1-8B-Instruct",
+    "meta-llama/Llama-4-Scout-17B-16E-Instruct",
+    "google/gemma-1.1-7b-it",
+    "Qwen/Qwen1.5-14B-Chat",
+    "mistralai/Mistral-7B-Instruct-v0.3"
+]
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# ----------------------
+# 📚 NLTK Setup
+# ----------------------
 try:
     nltk.data.find('tokenizers/punkt')
-except nltk.downloader.DownloadError:
-    nltk.download('punkt')
 except LookupError:
     nltk.download('punkt')
 
-# ---------------------------
-# 📄 Utility: Read PDF to text (with OCR fallback)
-# ---------------------------
-# This combines logic from extract_text_from_pdf and extract_text_from_page
-def extract_text_from_page_with_ocr(page):
+# ----------------------
+# 📄 Utility Functions
+# ----------------------
+def extract_text_or_ocr(page):
     text = page.get_text().strip()
     if text:
-        return text, False  # native text found, no OCR needed
+        return text, False
+    pix = page.get_pixmap(dpi=300)
+    img_data = pix.tobytes("png")
+    img = Image.open(io.BytesIO(img_data))
+    return pytesseract.image_to_string(img).strip(), True
 
-    # If native text is missing, try OCR
-    try:
-        pix = page.get_pixmap(dpi=300)
-        img_data = pix.tobytes("png")
-        img = Image.open(io.BytesIO(img_data))
-        ocr_text = pytesseract.image_to_string(img).strip()
-        return ocr_text, True
-    except Exception as e:
-        print(f"OCR failed for a page: {e}")
-        return "", False # Return empty and indicate OCR was not used if it fails
-
-
-# ---------------------------
-# 🧹 Clean up lines (from original notebook)
-# ---------------------------
 def clean_text(text):
-    lines = text.splitlines()
-    lines = [line.strip() for line in lines if line.strip()]
-    return "\n".join(lines)
+    return "\n".join([line.strip() for line in text.splitlines() if line.strip()])
 
-# ---------------------------
-# ✂️ Sentence Tokenizer (from original notebook)
-# ---------------------------
 def tokenize_sentences(text):
     return sent_tokenize(text)
 
-# ---------------------------
-# 📦 Chunk into fixed size blocks (from original notebook)
-# ---------------------------
-def split_into_chunks(sentences, max_tokens=750, overlap=100):
-    chunks = []
-    current_chunk = []
-    current_len = 0
-
+def split_chunks(sentences, max_tokens=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
+    chunks, chunk, length = [], [], 0
     for sentence in sentences:
-        token_count = len(sentence.split())
-        # Check if adding the next sentence exceeds max_tokens
-        # If it does, and the current chunk is not empty, save the current chunk
-        if current_len + token_count > max_tokens and current_chunk:
-            chunks.append(" ".join(current_chunk))
-            # Start the next chunk with the overlap
-            current_chunk = current_chunk[-overlap:]
-            # Recalculate current_len based on the overlap
-            current_len = sum(len(s.split()) for s in current_chunk)
-
-        # Add the current sentence and update length
-        current_chunk.append(sentence)
-        current_len += token_count
-
-    # Add the last chunk if it's not empty
-    if current_chunk:
-        chunks.append(" ".join(current_chunk))
-
+        count = len(sentence.split())
+        if length + count > max_tokens and chunk:
+            chunks.append(" ".join(chunk))
+            chunk = chunk[-overlap:]
+            length = sum(len(s.split()) for s in chunk)
+        chunk.append(sentence)
+        length += count
+    if chunk: chunks.append(" ".join(chunk))
     return chunks
 
-
-# ---------------------------
-# 🧠 Extract Metadata from Filename (from original notebook)
-# ---------------------------
-def extract_metadata_from_filename(filename):
+def extract_metadata(filename):
     name = filename.lower().replace("_", " ").replace("-", " ")
-
-    metadata = {
-        "model": "unknown",
-        "doc_type": "unknown",
-        "brand": "life fitness" # Assuming 'life fitness' is constant based on your notebook
-    }
-
-    if "om" in name or "owner" in name:
-        metadata["doc_type"] = "owner manual"
-    elif "sm" in name or "service" in name:
-        metadata["doc_type"] = "service manual"
-    elif "assembly" in name:
-        metadata["doc_type"] = "assembly instructions"
-    elif "alert" in name:
-        metadata["doc_type"] = "installer alert"
-    elif "parts" in name:
-        metadata["doc_type"] = "parts manual"
-    elif "bulletin" in name:
-        metadata["doc_type"] = "service bulletin"
-
-    known_models = [
-        "se3hd", "se3", "se4", "symbio", "explore", "integrity x", "integrity sl",
-        "everest", "engage", "inspire", "discover", "95t", "95x", "95c", "95r", "97c"
-    ]
-
-    for model in known_models:
-        # Use regex for more robust matching if needed, but simple 'in' check from notebook
-        if model.replace(" ", "") in name.replace(" ", ""):
-            metadata["model"] = model
-            break
-
-    return metadata
-
-# ---------------------------
-# 🚀 Step 1: Process PDFs, Extract Pages with OCR
-# ---------------------------
-def process_pdfs_for_pages(pdf_folder, output_jsonl):
-    print("Starting PDF processing and OCR...")
-    all_pages = []
-    if not os.path.exists(pdf_folder):
-        print(f"Error: PDF folder not found at {pdf_folder}")
-        return [] # Return empty list if folder doesn't exist
-
-    pdf_files = [f for f in os.listdir(pdf_folder) if f.lower().endswith(".pdf")]
-    if not pdf_files:
-        print(f"No PDF files found in {pdf_folder}")
-        return []
-
-    for pdf_file in tqdm(pdf_files, desc="Scanning PDFs"):
-        path = os.path.join(pdf_folder, pdf_file)
-        try:
-            doc = fitz.open(path)
-            for page_num, page in enumerate(doc, start=1):
-                text, used_ocr = extract_text_from_page_with_ocr(page)
-                if text: # Only save pages with extracted text
-                    all_pages.append({
-                        "source_file": pdf_file,
-                        "page": page_num,
-                        "text": text,
-                        "ocr_used": used_ocr
-                    })
-            doc.close() # Close the document
-        except Exception as e:
-            print(f"Error processing {pdf_file}: {e}")
-            continue # Skip to the next file
-
-    with open(output_jsonl, "w", encoding="utf-8") as f:
-        for page in all_pages:
-            json.dump(page, f)
-            f.write("\n")
-
-    print(f"✅ Saved {len(all_pages)} pages to {output_jsonl} (with OCR fallback)")
-    return all_pages # Return the list of pages
-
-# ---------------------------
-# 🚀 Step 2: Chunk the Pages
-# ---------------------------
-def chunk_pages(input_jsonl, output_jsonl, chunk_size, chunk_overlap):
-    print("Starting page chunking...")
-    all_chunks = []
-    if not os.path.exists(input_jsonl):
-        print(f"Error: Input JSONL file not found at {input_jsonl}. Run PDF processing first.")
-        return []
-
-    try:
-        with open(input_jsonl, "r", encoding="utf-8") as f:
-            # Count lines for tqdm progress bar
-            total_lines = sum(1 for _ in f)
-            f.seek(0) # Reset file pointer to the beginning
-
-            for line in tqdm(f, total=total_lines, desc="Chunking pages"):
-                try:
-                    page = json.loads(line)
-                    source_file = page["source_file"]
-                    page_number = page["page"]
-                    text = page["text"]
-
-                    metadata = extract_metadata_from_filename(source_file)
-                    sentences = tokenize_sentences(clean_text(text)) # Clean and tokenize the page text
-                    chunks = split_into_chunks(sentences, max_tokens=chunk_size, overlap=chunk_overlap)
-
-                    for i, chunk in enumerate(chunks):
-                         # Ensure chunk text is not empty
-                        if chunk.strip():
-                            all_chunks.append({
-                                "source_file": source_file,
-                                "chunk_id": f"{source_file}::page_{page_number}::chunk_{i+1}",
-                                "page": page_number,
-                                "ocr_used": page.get("ocr_used", False), # Use .get for safety
-                                "model": metadata.get("model", "unknown"),
-                                "doc_type": metadata.get("doc_type", "unknown"),
-                                "brand": metadata.get("brand", "life fitness"),
-                                "text": chunk.strip() # Ensure no leading/trailing whitespace
-                            })
-                except json.JSONDecodeError:
-                    print(f"Skipping invalid JSON line: {line}")
-                except Exception as e:
-                    print(f"Error processing page from {line}: {e}")
-                    continue # Continue with the next line
-
-    except Exception as e:
-        print(f"Error opening or reading input JSONL file: {e}")
-        return []
-
-
-    if not all_chunks:
-        print("No chunks were created.")
-
-    with open(output_jsonl, "w", encoding="utf-8") as f:
-        for chunk in all_chunks:
-            json.dump(chunk, f)
-            f.write("\n")
-
-    print(f"✅ Done! {len(all_chunks)} chunks saved to {output_jsonl}")
-    return all_chunks # Return the list of chunks
-
-# ---------------------------
-# 🚀 Step 3: Embed Chunks into Chroma
-# ---------------------------
-def embed_chunks_into_chroma(jsonl_path, chroma_path, collection_name):
-    print("Starting ChromaDB embedding...")
-    try:
-        embedder = SentenceTransformer("all-MiniLM-L6-v2")
-        embedder.eval()
-        print("✅ SentenceTransformer model loaded.")
-    except Exception as e:
-        print(f"❌ Error loading SentenceTransformer model: {e}")
-        return None, "Error loading SentenceTransformer model."
-
-    try:
-        # Use a persistent client
-        client = chromadb.PersistentClient(path=chroma_path)
-        # Check if collection exists and delete if it does to rebuild
-        try:
-            client.get_collection(name=collection_name)
-            client.delete_collection(collection_name)
-            print(f"Deleted existing collection: {collection_name}")
-        except Exception: # Collection does not exist, which is fine
-            pass
-        collection = client.create_collection(name=collection_name)
-        print(f"✅ ChromaDB collection '{collection_name}' created.")
-    except Exception as e:
-        print(f"❌ Error initializing ChromaDB: {e}")
-        return None, "Error initializing ChromaDB."
-
-    texts, metadatas, ids = [], [], []
-    batch_size = 16 # Define batch size for embedding
-
-    if not os.path.exists(jsonl_path):
-        print(f"Error: Input JSONL file not found at {jsonl_path}. Run chunking first.")
-        return None, "Input chunk file not found."
-
-    try:
-        with open(jsonl_path, "r", encoding="utf-8") as f:
-            # Count lines for tqdm progress bar
-            total_lines = sum(1 for _ in f)
-            f.seek(0) # Reset file pointer to the beginning
-
-            for line in tqdm(f, total=total_lines, desc="Embedding chunks"):
-                try:
-                    item = json.loads(line)
-                    texts.append(item.get("text", "")) # Use .get for safety
-                    ids.append(item.get("chunk_id", f"unknown_{len(ids)}")) # Ensure chunk_id exists
-                    # Prepare metadata, ensuring all keys are strings and handling potential missing keys
-                    metadata = {str(k): str(v) for k, v in item.items() if k != "text"}
-                    metadatas.append(metadata)
-
-                    if len(texts) >= batch_size:
-                        embeddings = embedder.encode(texts).tolist()
-                        collection.add(documents=texts, metadatas=metadatas, ids=ids, embeddings=embeddings)
-                        texts, metadatas, ids = [], [], [] # Reset batches
-
-                except json.JSONDecodeError:
-                    print(f"Skipping invalid JSON line during embedding: {line}")
-                except Exception as e:
-                    print(f"Error processing chunk line {line} during embedding: {e}")
-                    continue # Continue with the next line
-
-        # Add any remaining items in the last batch
-        if texts:
-            embeddings = embedder.encode(texts).tolist()
-            collection.add(documents=texts, metadatas=metadatas, ids=ids, embeddings=embeddings)
-
-        print("✅ All OCR-enhanced chunks embedded in Chroma!")
-        return collection, None # Return collection and no error
-
-    except Exception as e:
-        print(f"❌ Error reading input JSONL file for embedding: {e}")
-        return None, "Error reading input file for embedding."
-
-
-# ---------------------------
-# 🧠 Load Hugging Face Model and Tokenizer
-# ---------------------------
-# This needs to happen after imports but before the Gradio interface
-tokenizer = None
-model = None
-pipe = None
-
-print(f"Attempting to load Hugging Face model: {HF_MODEL_ID}")
-print(f"Using HF_TOKEN (present: {HF_TOKEN is not None})")
-
-if not HF_TOKEN:
-    print("❌ HF_TOKEN environment variable not set. Cannot load Hugging Face model.")
-else:
-    try:
-        # Check if CUDA is available
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Using device: {device}")
-
-        # Load tokenizer and model
-        tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID, token=HF_TOKEN)
-        model = AutoModelForCausalLM.from_pretrained(
-            HF_MODEL_ID,
-            token=HF_TOKEN,
-            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, # Use bfloat16 on GPU
-            device_map="auto" if torch.cuda.is_available() else None # Auto device mapping on GPU
-        ).to(device) # Move model to selected device
-
-        # Create a pipeline for easy inference
-        pipe = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            max_new_tokens=512,
-            temperature=0.1,
-            top_p=0.9,
-            do_sample=True,
-            device=0 if torch.cuda.is_available() else -1 # Specify device for pipeline
-        )
-
-        print(f"✅ Successfully loaded Hugging Face model: {HF_MODEL_ID} on {device}")
-
-    except Exception as e:
-        print(f"❌ Error loading Hugging Face model: {e}")
-        print("Please ensure:")
-        print("- The HF_TOKEN secret is set in your Hugging Face Space settings.")
-        print("- Your Space has sufficient resources (GPU, RAM) for the model.")
-        print("- You have accepted the model's terms on Hugging Face (if required).")
-        tokenizer, model, pipe = None, None, None # Set to None if loading fails
-
-
-# ---------------------------
-# 🔎 Query Function (Uses Embedder and Chroma)
-# ---------------------------
-# Embedder is loaded during the embedding step, need to ensure it's accessible
-embedder = None # Initialize embedder as None
-
-def query_manuals(question, model_filter=None, doc_type_filter=None, top_k=5, rerank_keywords=None):
-    global embedder # Access the global embedder variable
-
-    if collection is None or embedder is None:
-        print("⚠️ ChromaDB or Embedder not loaded. Cannot perform vector search.")
-        return [] # Return empty if Chroma or Embedder is not loaded
-
-    where_filter = {}
-    if model_filter:
-        where_filter["model"] = model_filter.lower()
-    if doc_type_filter:
-        where_filter["doc_type"] = doc_type_filter.lower()
-
-    # ChromaDB query expects a dictionary for 'where'
-    results = collection.query(
-        query_texts=[question],
-        n_results=top_k * 5,  # fetch more for reranking
-        where={} if not where_filter else where_filter # Pass empty dict if no filter
-    )
-
-
-    if not results or not results.get("documents") or not results["documents"][0]:
-        return []  # No matches
-
-    try:
-        question_embedding = embedder.encode(question, convert_to_tensor=True)
-    except Exception as e:
-        print(f"Error encoding question: {e}")
-        return [] # Return empty if embedding fails
-
-    # Step 3: Compute semantic + keyword score
-    reranked = []
-    # Ensure results["documents"] and results["metadatas"] are not empty before iterating
-    if results.get("documents") and results["documents"][0]:
-        for i, text in enumerate(results["documents"][0]):
-            meta = results["metadatas"][0][i]
-
-            # Handle potential encoding errors during text embedding
-            try:
-                embedding = embedder.encode(text, convert_to_tensor=True)
-                 # Semantic similarity
-                similarity_score = float(util.cos_sim(question_embedding, embedding))
-            except Exception as e:
-                 print(f"Error encoding chunk text for reranking: {e}. Skipping chunk.")
-                 continue # Skip this chunk if encoding fails
-
-
-            # Keyword score
-            keyword_score = 0
-            if rerank_keywords and text: # Ensure text is not None or empty
-                for kw in rerank_keywords:
-                    if kw.lower() in text.lower():
-                        keyword_score += 1
-
-            # Combine with tunable weights
-            # Weights should sum to 1 for a simple weighted average
-            final_score = (0.8 * similarity_score) + (0.2 * keyword_score)
-
-            reranked.append({
-                "score": final_score,
-                "text": text,
-                "metadata": meta
-            })
-
-    # Sort by combined score
-    reranked.sort(key=lambda x: x["score"], reverse=True)
-    return reranked[:top_k]
-
-
-# ---------------------------
-# 💬 Ask Hugging Face Model
-# ---------------------------
-def ask_hf_model(prompt):
-    if pipe is None:
-        return "Hugging Face model not loaded. Cannot generate response."
-    try:
-        # Use the Llama 3.1 chat template
-        messages = [
-            {"role": "system", "content": "You are a technical assistant trained to answer questions using equipment manuals. Use only the provided context to answer the question. If the answer is not clearly in the context, reply: 'I don't know.'"},
-            {"role": "user", "content": prompt}
-        ]
-
-        # Apply chat template and generate text
-        prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-
-        outputs = pipe(
-            prompt_text,
-            do_sample=True,
-            temperature=0.1, # Keep temperature low for more factual answers
-            top_p=0.9,
-            max_new_tokens=512,
-            pad_token_id=tokenizer.eos_token_id # Set pad_token_id for generation
-        )
-        # The output includes the prompt, we need to extract just the generated part
-        # Find the end of the prompt text in the generated output
-        generated_text = outputs[0]["generated_text"]
-        # The chat template adds the assistant's turn token, look for that to find the response start
-        response_start_token = tokenizer.apply_chat_template([{"role": "assistant", "content": ""}], tokenize=False, add_generation_prompt=False)
-        response_start_index = generated_text.find(response_start_token)
-
-        if response_start_index != -1:
-             response = generated_text[response_start_index + len(response_start_token):].strip()
-        else:
-             # Fallback if the assistant token isn't found
-             response = generated_text.strip()
-
-        # Remove any trailing EOS tokens or similar artifacts
-        if response.endswith(tokenizer.eos_token):
-            response = response[:-len(tokenizer.eos_token)].strip()
-
-
-        return response
-
-    except Exception as e:
-        return f"❌ Error generating response from Hugging Face model: {str(e)}"
-
-# ---------------------------
-# 🎯 Full RAG Pipeline
-# ---------------------------
-def run_rag_qa(user_question, model_filter=None, doc_type_filter=None): # Added filters as optional inputs
-    # Ensure ChromaDB and the HF model pipeline are loaded before proceeding
-    if collection is None:
-        return "ChromaDB is not loaded. Ensure PDFs are in ./Manuals and the app started correctly."
-    if pipe is None:
-        return "Hugging Face model pipeline is not loaded. Ensure HF_TOKEN is set and the model loaded successfully."
-
-    results = query_manuals(
-        question=user_question,
-        model_filter=model_filter, # Use the optional filter inputs
-        doc_type_filter=doc_type_filter,
-        top_k=MAX_CONTEXT_CHUNKS,
-        rerank_keywords=["diagnostic", "immobilize", "system", "screen", "service", "error"] # Example keywords
-    )
-
-    if not results:
-        # Attempt a broader search if initial filter yields no results
-        if model_filter or doc_type_filter:
-             print("No results with specified filters, trying broader search...")
-             results = query_manuals(
-                question=user_question,
-                model_filter=None, # Remove filters for broader search
-                doc_type_filter=None,
-                top_k=MAX_CONTEXT_CHUNKS,
-                rerank_keywords=["diagnostic", "immobilize", "system", "screen", "service", "error"]
-            )
-             if not results:
-                return "No relevant documents found for the query, even with broader search."
-        else:
-            return "No relevant documents found for the query."
-
-
-    context = "\n\n".join([f"Source File: {r['metadata'].get('source_file', 'N/A')}, Page: {r['metadata'].get('page', 'N/A')}\nText: {r['text'].strip()}" for r in results])
-
-    prompt = f"""
-Context:
-{context}
-
-Question: {user_question}
-"""
-
-    return ask_hf_model(prompt)
-
-# ---------------------------
-# --- Initial Setup ---
-# This code runs when the app starts on Hugging Face Spaces
-# It processes PDFs, chunks, and builds the ChromaDB
-# ---------------------------
-
-print("Starting initial setup...")
-
-# Ensure Tesseract is available on the system (Hugging Face Spaces usually has it, but this command is good practice)
-# Using ! in app.py is generally discouraged, better to ensure the environment has it
-# For HF Spaces, you might need to use a Dockerfile or rely on the default environment.
-# If Tesseract isn't found, the OCR part might fail.
-
-# Process PDFs and extract pages
-all_pages = process_pdfs_for_pages(pdf_folder, output_jsonl_pages)
-
-# Chunk the pages
-all_chunks = []
-if all_pages: # Only chunk if pages were processed
-    all_chunks = chunk_pages(output_jsonl_pages, output_jsonl_chunks, chunk_size, chunk_overlap)
-
-# Embed chunks into ChromaDB
-collection = None # Initialize collection
-if all_chunks: # Only embed if chunks were created
-    collection, embed_error = embed_chunks_into_chroma(output_jsonl_chunks, chroma_path, collection_name)
-    if embed_error:
-        print(f"Error during embedding: {embed_error}")
-
-
-print("Initial setup complete.")
-
-# ---------------------------
-# 🖥️ Gradio Interface
-# ---------------------------
-# Only define and launch the interface if the necessary components loaded
-if collection is not None and pipe is not None:
-    with gr.Blocks() as demo:
-        gr.Markdown("""# 🧠 Manual QA via Hugging Face Llama 3.1
-    Ask a technical question and get answers using your own PDF manual database and a Hugging Face model.
-    **Note:** Initial startup might take time to process manuals and build the search index. Ensure your `Manuals` folder is uploaded and the `HF_TOKEN` secret is set in Space settings.
-    """)
-        with gr.Row():
-            question = gr.Textbox(label="Your Question", placeholder="e.g. How do I access diagnostics on the SE3 console?")
-        with gr.Row():
-             model_filter_input = gr.Textbox(label="Filter by Model (Optional)", placeholder="e.g. se3hd")
-             doc_type_filter_input = gr.Dropdown(label="Filter by Document Type (Optional)", choices=["owner manual", "service manual", "assembly instructions", "installer alert", "parts manual", "service bulletin", "unknown", None], value=None, allow_custom_value=True)
-
-        submit = gr.Button("🔍 Ask")
-        answer = gr.Textbox(label="Answer", lines=10) # Increased lines for better readability
-
-        # Call the run_rag_qa function when the button is clicked
-        submit.click(
-            fn=run_rag_qa,
-            inputs=[question, model_filter_input, doc_type_filter_input],
-            outputs=[answer]
-        )
-
-    # In Hugging Face Spaces, the app is launched automatically.
-    # The demo.launch() call is removed.
-    # demo.launch()
-else:
-    print("Gradio demo will not launch because RAG components (ChromaDB or HF Model) failed to load during setup.")
-    # You could add a simple Gradio interface here to show an error message
-    # if you wanted to provide user feedback in the Space UI even on failure.
-    # Example:
-    # with gr.Blocks() as error_demo:
-    #     gr.Markdown("## Application Failed to Load")
-    #     gr.Textbox(label="Error Details", value="RAG components (ChromaDB or HF Model) failed to initialize. Check logs and Space settings (HF_TOKEN, resources).", interactive=False)
-    # error_demo.launch()
+    meta = {"model": "unknown", "doc_type": "unknown", "brand": "life fitness"}
+    if "om" in name or "owner" in name: meta["doc_type"] = "owner manual"
+    elif "sm" in name or "service" in name: meta["doc_type"] = "service manual"
+    elif "assembly" in name: meta["doc_type"] = "assembly instructions"
+    elif "alert" in name: meta["doc_type"] = "installer alert"
+    elif "parts" in name: meta["doc_type"] = "parts manual"
+    elif "bulletin" in name: meta["doc_type"] = "service bulletin"
+    for kw in ["se3hd", "se3", "se4", "symbio", "explore", "integrity x", "integrity sl", "everest", "engage", "inspire", "discover", "95t", "95x", "95c", "95r", "97c"]:
+        if kw.replace(" ", "") in name.replace(" ", ""): meta["model"] = kw
+    return meta
+
+# ----------------------
+# 🧠 Load LLM
+# ----------------------
+def load_llm(model_id):
+    tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
+    model = AutoModelForCausalLM.from_pretrained(model_id, token=HF_TOKEN, torch_dtype=torch.float32)
+    return pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
+
+# ----------------------
+# 🧠 Chroma + Embed
+# ----------------------
+def embed_pdfs():
+    os.makedirs(CHROMA_PATH, exist_ok=True)
+    client = chromadb.PersistentClient(path=CHROMA_PATH)
+    if COLLECTION_NAME in [c.name for c in client.list_collections()]:
+        client.delete_collection(COLLECTION_NAME)
+    collection = client.create_collection(COLLECTION_NAME)
+    embedder = SentenceTransformer("all-MiniLM-L6-v2")
+
+    for file in tqdm(os.listdir(PDF_DIR)):
+        if not file.lower().endswith(".pdf"): continue
+        doc = fitz.open(os.path.join(PDF_DIR, file))
+        meta = extract_metadata(file)
+        for page_num, page in enumerate(doc, 1):
+            text, _ = extract_text_or_ocr(page)
+            if not text.strip(): continue
+            sents = tokenize_sentences(clean_text(text))
+            chunks = split_chunks(sents)
+            for i, chunk in enumerate(chunks):
+                chunk_id = f"{file}::p{page_num}::c{i}"
+                emb = embedder.encode([chunk])[0].tolist()
+                collection.add(
+                    documents=[chunk],
+                    ids=[chunk_id],
+                    embeddings=[emb],
+                    metadatas=[{**meta, "source_file": file, "page": page_num}]
+                )
+    return collection, embedder
+
+# ----------------------
+# 🔍 RAG Pipeline
+# ----------------------
+def answer_query(q, model_id):
+    collection, embedder = embed_pdfs()
+    pipe = load_llm(model_id)
+    emb_q = embedder.encode([q])[0].tolist()
+    results = collection.query(query_embeddings=[emb_q], n_results=MAX_CONTEXT_CHUNKS)
+    context = "\n\n".join(results['documents'][0])
+    prompt = f"Use the context below to answer the question.\nContext:\n{context}\n\nQuestion: {q}\nAnswer:"
+    return pipe(prompt)[0]['generated_text'].split("Answer:")[-1].strip()
+
+# ----------------------
+# 🚀 Gradio UI
+# ----------------------
+with gr.Blocks() as app:
+    gr.Markdown("""# SmartManuals-AI
+**Local-first document QA** powered by OCR, ChromaDB & your choice of LLM (via Hugging Face).
+""")
+    with gr.Row():
+        question = gr.Textbox(placeholder="Ask a question from the manuals...", label="Question")
+        model_choice = gr.Dropdown(label="Choose Model", choices=MODEL_OPTIONS, value=MODEL_OPTIONS[0])
+    output = gr.Textbox(label="Answer", lines=10)
+    run = gr.Button("Run RAG")
+    run.click(fn=answer_query, inputs=[question, model_choice], outputs=output)
+
+if __name__ == "__main__":
+    app.launch()