Update app.py

app.py

@@ -1,97 +1,590 @@
-# app.py
-from flask import Flask, send_from_directory, jsonify, request
 import os
+import io
+import requests
+import logging
+import re
 import json
-import errno
+import base64
+from flask import Flask, request, render_template, jsonify, send_file, Response
+from PyPDF2 import PdfReader, PdfWriter
+import pytesseract
+from pdf2image import convert_from_bytes
+from PIL import Image
+from datasets import Dataset, load_dataset
+from sentence_transformers import SentenceTransformer
+from datetime import datetime
+from numpy import dot
+from numpy.linalg import norm
+from huggingface_hub import HfApi, hf_hub_download
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+import torch
+import chromadb
+from chromadb.utils import embedding_functions
+import shutil
+
+# Set up logging
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Set cache, uploads, and output directories
+os.environ["HF_HOME"] = "/app/cache"
+os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
+os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/app/cache"
+os.environ["XDG_CACHE_HOME"] = "/app/cache"
+UPLOADS_DIR = "/app/uploads"
+PAGES_DIR = os.path.join(UPLOADS_DIR, "pages")
+OUTPUT_DIR = "/app/output"
+COMBINED_PDF_PATH = os.path.join(OUTPUT_DIR, "combined_output.pdf")
+PROGRESS_JSON_PATH = os.path.join(OUTPUT_DIR, "progress_log.json")
+CHROMA_DB_PATH = os.path.join(OUTPUT_DIR, "chromadb")
+os.makedirs(PAGES_DIR, exist_ok=True)
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 app = Flask(__name__)
 
-# Directory to store settings
-SETTINGS_DIR = "settings"
+# Hugging Face Hub configuration
+HF_TOKEN = os.getenv("HF_TOKEN")
+HF_DATASET_REPO = "broadfield-dev/pdf-ocr-dataset"
+HF_API = HfApi()
+
+# Tracking file for resuming
+TRACKING_FILE = "/app/cache/processing_state.json"
+
+# Load sentence transformer
+try:
+    embedder = SentenceTransformer('all-MiniLM-L6-v2', cache_folder="/app/cache")
+    logger.info("SentenceTransformer loaded successfully")
+except Exception as e:
+    logger.error(f"Failed to load SentenceTransformer: {e}")
 
-# Ensure the settings directory exists and has the correct permissions
+# Initialize TrOCR (CPU-only)
 try:
-    if not os.path.exists(SETTINGS_DIR):
-        os.makedirs(SETTINGS_DIR, mode=0o775)
-    os.chmod(SETTINGS_DIR, 0o775)
-except OSError as e:
-    print(f"Error setting up settings directory: {e}")
-
-# Real-world simulation parameters (unscaled)
-simulation_params = {
-    "sun": {
-        "mass": 1.989e30,  # Real mass in kg
-        "position": [0, 0, 0],
-        "orbital_velocity": 0,  # Sun is stationary
-    },
-    "earth": {
-        "mass": 5.972e24,  # Real mass in kg
-        "position": [149.6e6, 0, 0],  # 1 AU in km
-        "orbital_velocity": 29.8,  # Real orbital velocity in km/s
-    },
-    "mars": {
-        "mass": 6.417e23,  # Real mass in kg
-        "position": [227.9e6, 0, 0],  # 1.52 AU in km
-        "orbital_velocity": 24.1,  # Real orbital velocity in km/s
-    },
-    "fluid_speed": 0.1,
-    "fluid_friction": 0.9,
-    "fluid_deflection": 0.1,
-}
-
-# Serve the frontend
-@app.route('/')
-def serve_index():
-    return send_from_directory('static', 'index.html')
-
-# Serve static files (CSS, JS)
-@app.route('/static/<path:path>')
-def serve_static(path):
-    return send_from_directory('static', path)
-
-# API to get simulation parameters
-@app.route('/api/params', methods=['GET'])
-def get_params():
-    return jsonify(simulation_params)
-
-# API to update simulation parameters
-@app.route('/api/params', methods=['POST'])
-def update_params():
-    global simulation_params
-    data = request.get_json()
-    simulation_params.update(data)
-    return jsonify({"status": "success", "params": simulation_params})
-
-# API to save settings to a JSON file
-@app.route('/api/save', methods=['POST'])
-def save_settings():
+    trocr_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
+    trocr_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
+    trocr_model.to("cpu").eval()
+    logger.info("TrOCR initialized successfully on CPU")
+except Exception as e:
+    logger.error(f"Failed to initialize TrOCR: {e}")
+    trocr_model = None
+    trocr_processor = None
+
+# Initialize ChromaDB
+chroma_client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
+sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
+chroma_collection = chroma_client.get_or_create_collection(name="pdf_pages", embedding_function=sentence_transformer_ef)
+
+# Load or initialize progress log
+def load_progress_log(storage_mode):
+    if storage_mode == "hf":
+        try:
+            progress_file = hf_hub_download(repo_id=HF_DATASET_REPO, filename="progress_log.json", repo_type="dataset", token=HF_TOKEN)
+            with open(progress_file, "r") as f:
+                return json.load(f)
+        except Exception as e:
+            logger.info(f"No HF progress log found or error loading: {e}, initializing new log")
+            return {"urls": {}}
+    else:  # local
+        if os.path.exists(PROGRESS_JSON_PATH):
+            with open(PROGRESS_JSON_PATH, "r") as f:
+                return json.load(f)
+        return {"urls": {}}
+
+def save_progress_log(progress_log, storage_mode):
+    if storage_mode == "hf":
+        with open("/app/cache/progress_log.json", "w") as f:
+            json.dump(progress_log, f)
+        HF_API.upload_file(
+            path_or_fileobj="/app/cache/progress_log.json",
+            path_in_repo="progress_log.json",
+            repo_id=HF_DATASET_REPO,
+            repo_type="dataset",
+            token=HF_TOKEN
+        )
+        logger.info("Progress log updated in Hugging Face dataset")
+    else:  # local
+        with open(PROGRESS_JSON_PATH, "w") as f:
+            json.dump(progress_log, f)
+        logger.info("Progress log updated locally")
+
+# Tesseract OCR with bounding boxes
+def ocr_with_tesseract(pdf_bytes, page_num):
+    try:
+        images = convert_from_bytes(pdf_bytes, first_page=page_num+1, last_page=page_num+1)
+        if not images:
+            logger.info(f"Page {page_num + 1} is blank")
+            return {"page_num": page_num + 1, "text": "Blank page", "words": []}
+        image = images[0]
+        data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
+        text = pytesseract.image_to_string(image)
+        words = [
+            {"text": data["text"][i], "left": data["left"][i], "top": data["top"][i], "width": data["width"][i], "height": data["height"][i]}
+            for i in range(len(data["text"])) if data["text"][i].strip()
+        ]
+        logger.info(f"Tesseract processed page {page_num + 1} with {len(words)} words")
+        return {"page_num": page_num + 1, "text": text, "words": words}
+    except Exception as e:
+        logger.error(f"Tesseract error on page {page_num + 1}: {e}")
+        return {"page_num": page_num + 1, "text": f"Tesseract Error: {str(e)}", "words": []}
+
+# TrOCR OCR
+def ocr_with_trocr(pdf_bytes, page_num):
+    if not trocr_model or not trocr_processor:
+        logger.warning(f"TrOCR not available for page {page_num + 1}")
+        return {"page_num": page_num + 1, "text": "TrOCR not initialized", "words": []}
+    try:
+        images = convert_from_bytes(pdf_bytes, first_page=page_num+1, last_page=page_num+1)
+        if not images:
+            logger.info(f"Page {page_num + 1} is blank")
+            return {"page_num": page_num + 1, "text": "Blank page", "words": []}
+        image = images[0].convert("RGB")
+        pixel_values = trocr_processor(image, return_tensors="pt").pixel_values.to("cpu")
+        with torch.no_grad():
+            generated_ids = trocr_model.generate(pixel_values, max_length=50)
+        text = trocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        words = [{"text": word, "left": 0, "top": 0, "width": 0, "height": 0} for word in text.split()]
+        logger.info(f"TrOCR processed page {page_num + 1} with text: {text}")
+        return {"page_num": page_num + 1, "text": text, "words": words}
+    except Exception as e:
+        logger.error(f"TrOCR error on page {page_num + 1}: {e}")
+        return {"page_num": page_num + 1, "text": f"TrOCR Error: {str(e)}", "words": []}
+
+# Map Tesseract bounding boxes to OCR text
+def map_tesseract_to_ocr(tesseract_result, ocr_result):
+    if not tesseract_result["words"] or "Error" in ocr_result["text"]:
+        logger.info(f"Mapping skipped for page {tesseract_result['page_num']}: No Tesseract words or OCR error")
+        return {**ocr_result, "words": tesseract_result["words"]}
+    
+    ocr_text = ocr_result["text"]
+    tesseract_words = tesseract_result["words"]
+    
+    sentences = re.split(r'(?<=[.!?])\s+', ocr_text.strip())
+    sentence_embeddings = embedder.encode(sentences)
+    
+    mapped_words = []
+    for word in tesseract_words:
+        word_embedding = embedder.encode(word["text"])
+        similarities = [
+            dot(word_embedding, sent_emb) / (norm(word_embedding) * norm(sent_emb)) if norm(sent_emb) != 0 else 0
+            for sent_emb in sentence_embeddings
+        ]
+        best_match_idx = similarities.index(max(similarities))
+        best_sentence = sentences[best_match_idx]
+        if word["text"].lower() in best_sentence.lower():
+            mapped_words.append(word)
+        else:
+            mapped_words.append(word)
+    logger.info(f"Mapped {len(mapped_words)} words for page {tesseract_result['page_num']}")
+    return {**ocr_result, "words": mapped_words}
+
+# Update combined PDF
+def update_combined_pdf(pdf_bytes, page_num):
+    pdf_reader = PdfReader(io.BytesIO(pdf_bytes))
+    page = pdf_reader.pages[page_num]
+    writer = PdfWriter()
+    if os.path.exists(COMBINED_PDF_PATH):
+        existing_pdf = PdfReader(COMBINED_PDF_PATH)
+        for p in existing_pdf.pages:
+            writer.add_page(p)
+    writer.add_page(page)
+    with open(COMBINED_PDF_PATH, "wb") as f:
+        writer.write(f)
+    logger.info(f"Updated combined PDF with page {page_num + 1}")
+
+# Process page
+def process_page(pdf_bytes, page_num, ocr_backend, filename, tracking_state, storage_mode):
+    tesseract_result = ocr_with_tesseract(pdf_bytes, page_num)
+    ocr_result = ocr_with_trocr(pdf_bytes, page_num) if ocr_backend == "trocr" else ocr_with_tesseract(pdf_bytes, page_num)
+    combined_result = map_tesseract_to_ocr(tesseract_result, ocr_result)
+    
+    local_page_path = os.path.join(PAGES_DIR, f"{filename}_page_{combined_result['page_num']}_{datetime.now().strftime('%Y%m%d%H%M%S')}.pdf")
+    writer = PdfWriter()
+    pdf_reader = PdfReader(io.BytesIO(pdf_bytes))
+    writer.add_page(pdf_reader.pages[page_num])
+    with open(local_page_path, "wb") as f:
+        writer.write(f)
+    
+    if storage_mode == "hf":
+        remote_page_path = f"pages/{os.path.basename(local_page_path)}"
+        HF_API.upload_file(
+            path_or_fileobj=local_page_path,
+            path_in_repo=remote_page_path,
+            repo_id=HF_DATASET_REPO,
+            repo_type="dataset",
+            token=HF_TOKEN
+        )
+        logger.info(f"Uploaded page to {HF_DATASET_REPO}/{remote_page_path}")
+        combined_result["page_file"] = remote_page_path
+    else:  # local
+        update_combined_pdf(pdf_bytes, page_num)
+        combined_result["page_file"] = local_page_path
+    
+    combined_result["pdf_page"] = tracking_state["last_offset"] + page_num
+    
+    # Update ChromaDB
+    chroma_collection.add(
+        documents=[combined_result["text"]],
+        metadatas=[{"filename": filename, "page_num": combined_result["page_num"], "page_file": combined_result["page_file"], "words": json.dumps(combined_result["words"])}],
+        ids=[f"{filename}_page_{combined_result['page_num']}"]
+    )
+    logger.info(f"Added page {combined_result['page_num']} to ChromaDB")
+    
+    return combined_result
+
+# Extract PDF URLs from text
+def extract_pdf_urls(text):
+    url_pattern = r'(https?://[^\s]+?\.pdf)'
+    return re.findall(url_pattern, text)
+
+# Load or initialize tracking state
+def load_tracking_state():
+    if os.path.exists(TRACKING_FILE):
+        with open(TRACKING_FILE, "r") as f:
+            return json.load(f)
+    return {"processed_urls": {}, "last_offset": 0}
+
+def save_tracking_state(state):
+    with open(TRACKING_FILE, "w") as f:
+        json.dump(state, f)
+
+# Push to Hugging Face Dataset
+def push_to_hf_dataset(new_data):
     try:
-        filename = os.path.join(SETTINGS_DIR, "settings.json")
-        with open(filename, 'w') as f:
-            json.dump(simulation_params, f, indent=4)
-        return jsonify({"status": "success", "message": "Settings saved successfully"})
-    except PermissionError as e:
-        return jsonify({"status": "error", "message": "Permission denied: Unable to save settings. Please check directory permissions."}), 500
+        for item in new_data:
+            if "url" not in item or not isinstance(item["url"], str):
+                logger.error(f"Invalid item in new_data: {item}")
+                raise ValueError(f"Each item must have a valid 'url' key; found {item}")
+        
+        try:
+            dataset = load_dataset(HF_DATASET_REPO, token=HF_TOKEN, cache_dir="/app/cache")
+            existing_data = dataset["train"].to_dict()
+            logger.info(f"Loaded existing dataset with keys: {list(existing_data.keys())}")
+        except Exception as e:
+            logger.info(f"No existing dataset found or error loading: {e}, initializing new dataset")
+            existing_data = {"filename": [], "pages": [], "url": [], "embedding": [], "processed_at": [], "pdf_page_offset": []}
+        
+        required_keys = ["filename", "pages", "url", "embedding", "processed_at", "pdf_page_offset"]
+        for key in required_keys:
+            if key not in existing_data:
+                existing_data[key] = []
+                logger.warning(f"Initialized missing key '{key}' in existing_data")
+        
+        existing_urls = set(existing_data["url"])
+        for item in new_data:
+            logger.debug(f"Processing item: {item}")
+            if item["url"] not in existing_urls:
+                for key in required_keys:
+                    existing_data[key].append(item.get(key, None))
+                existing_urls.add(item["url"])
+                logger.info(f"Added new URL: {item['url']}")
+            else:
+                idx = existing_data["url"].index(item["url"])
+                existing_data["pages"][idx].extend(item["pages"])
+                existing_data["embedding"][idx] = item["embedding"]
+                existing_data["processed_at"][idx] = item["processed_at"]
+                logger.info(f"Updated existing URL: {item['url']}")
+        
+        updated_dataset = Dataset.from_dict(existing_data)
+        updated_dataset.push_to_hub(HF_DATASET_REPO, token=HF_TOKEN)
+        logger.info(f"Successfully appended/updated {len(new_data)} records to {HF_DATASET_REPO}")
     except Exception as e:
-        return jsonify({"status": "error", "message": f"Error saving settings: {str(e)}"}), 500
+        logger.error(f"Failed to push to HF Dataset: {str(e)}")
+        raise
+
+# Check if URL is fully processed
+def is_url_fully_processed(url, progress_log, total_pages):
+    return url in progress_log["urls"] and progress_log["urls"][url]["status"] == "completed" and progress_log["urls"][url]["processed_pages"] >= total_pages
 
-# API to load settings from a JSON file
-@app.route('/api/load', methods=['GET'])
-def load_settings():
-    global simulation_params
+# Process PDF URL with SSE
+def process_pdf_url(url, ocr_backend, tracking_state, progress_log, storage_mode):
+    filename = url.split("/")[-1]
     try:
-        filename = os.path.join(SETTINGS_DIR, "settings.json")
-        if os.path.exists(filename):
-            with open(filename, 'r') as f:
-                simulation_params = json.load(f)
-            return jsonify({"status": "success", "params": simulation_params})
+        yield f"data: {json.dumps({'status': 'fetching', 'filename': filename})}\n\n"
+        logger.info(f"Fetching PDF from {url}")
+        response = requests.get(url, timeout=10)
+        response.raise_for_status()
+        pdf_bytes = response.content
+        pdf_reader = PdfReader(io.BytesIO(pdf_bytes))
+        total_pages = len(pdf_reader.pages)
+        
+        progress_log["urls"].setdefault(url, {"status": "pending", "processed_pages": 0})
+        start_page = progress_log["urls"][url]["processed_pages"]
+        
+        if is_url_fully_processed(url, progress_log, total_pages):
+            yield f"data: {json.dumps({'status': 'skipped', 'filename': filename, 'message': 'URL already fully processed'})}\n\n"
+            return
+        
+        pages = []
+        for page_num in range(start_page, total_pages):
+            yield f"data: {json.dumps({'status': 'processing', 'filename': filename, 'page_num': page_num + 1, 'total_pages': total_pages})}\n\n"
+            page = process_page(pdf_bytes, page_num, ocr_backend, filename, tracking_state, storage_mode)
+            pages.append(page)
+            yield f"data: {json.dumps({'filename': filename, 'page': page})}\n\n"
+            progress_log["urls"][url]["processed_pages"] = page_num + 1
+            save_progress_log(progress_log, storage_mode)
+        
+        full_text = "\n\n".join(f"Page {page['page_num']}\n{page['text']}" for page in pages)
+        embedding = embedder.encode(full_text).tolist() if full_text.strip() else None
+        result = {
+            "filename": filename,
+            "pages": pages,
+            "url": url,
+            "embedding": embedding,
+            "processed_at": datetime.now().isoformat(),
+            "pdf_page_offset": tracking_state["last_offset"]
+        }
+        if storage_mode == "hf":
+            push_to_hf_dataset([result])
+        tracking_state["last_offset"] += total_pages - start_page
+        progress_log["urls"][url]["status"] = "completed"
+        save_tracking_state(tracking_state)
+        save_progress_log(progress_log, storage_mode)
+        yield f"data: {json.dumps({'status': 'completed', 'filename': filename, 'new_offset': tracking_state['last_offset']})}\n\n"
+        logger.info(f"Completed processing {filename} with new offset {tracking_state['last_offset']}")
+    except requests.RequestException as e:
+        logger.error(f"Failed to fetch PDF from {url}: {e}")
+        yield f"data: {json.dumps({'status': 'error', 'filename': filename, 'message': f'Error fetching PDF: {str(e)}'})}\n\n"
+    except Exception as e:
+        logger.error(f"Error processing {url}: {e}")
+        yield f"data: {json.dumps({'status': 'error', 'filename': filename, 'message': f'Error: {str(e)}'})}\n\n"
+
+# Process text content with SSE
+def process_text_content(text, filename, ocr_backend, tracking_state, progress_log, storage_mode):
+    try:
+        pdf_urls = extract_pdf_urls(text)
+        processed_urls = [url for url in pdf_urls if url in progress_log["urls"] and progress_log["urls"][url]["status"] == "completed"]
+        new_urls = [url for url in pdf_urls if url not in progress_log["urls"] or progress_log["urls"][url]["status"] != "completed"]
+        
+        initial_text = (f"Found {len(pdf_urls)} PDF URLs:\n" +
+                       f"Already processed: {len(processed_urls)}\n" + "\n".join(processed_urls) + "\n" +
+                       f"To process: {len(new_urls)}\n" + "\n".join(new_urls) + "\n\nProcessing...")
+        yield f"data: {json.dumps({'status': 'info', 'filename': filename, 'message': initial_text})}\n\n"
+        
+        for url in new_urls:
+            logger.info(f"Starting processing of {url} with offset {tracking_state['last_offset']}")
+            for event in process_pdf_url(url, ocr_backend, tracking_state, progress_log, storage_mode):
+                yield event
+    except Exception as e:
+        logger.error(f"Error processing text content for {filename}: {e}")
+        yield f"data: {json.dumps({'status': 'error', 'filename': filename, 'message': f'Error: {str(e)}'})}\n\n"
+
+# Home route
+@app.route("/", methods=["GET"])
+def index():
+    return render_template("index.html")
+
+# Process URL endpoint with GET
+@app.route("/process_url", methods=["GET"])
+def process_url():
+    url = request.args.get("url")
+    ocr_backend = request.args.get("ocr_backend", "trocr")
+    storage_mode = request.args.get("storage_mode", "hf")
+    
+    if not url:
+        return jsonify({"error": "No URL provided"}), 400
+    
+    tracking_state = load_tracking_state()
+    progress_log = load_progress_log(storage_mode)
+    
+    def generate():
+        logger.info(f"Processing URL: {url} with ocr_backend={ocr_backend}, storage_mode={storage_mode}, starting offset={tracking_state['last_offset']}")
+        if url.endswith(".pdf"):
+            for event in process_pdf_url(url, ocr_backend, tracking_state, progress_log, storage_mode):
+                yield event
+        elif url.endswith(".txt"):
+            try:
+                response = requests.get(url, timeout=10)
+                response.raise_for_status()
+                text = response.text
+                filename = url.split("/")[-1]
+                logger.info(f"Fetched text from {url}")
+                for event in process_text_content(text, filename, ocr_backend, tracking_state, progress_log, storage_mode):
+                    yield event
+            except requests.RequestException as e:
+                logger.error(f"Failed to fetch text from {url}: {e}")
+                yield f"data: {json.dumps({'status': 'error', 'filename': url, 'message': f'Error fetching URL: {str(e)}'})}\n\n"
         else:
-            return jsonify({"status": "error", "message": "No saved settings found"}), 404
-    except PermissionError as e:
-        return jsonify({"status": "error", "message": "Permission denied: Unable to load settings. Please check directory permissions."}), 500
+            yield f"data: {json.dumps({'status': 'error', 'filename': url, 'message': 'Unsupported URL format. Must end in .pdf or .txt'})}\n\n"
+        logger.info(f"Finished processing URL: {url}")
+    
+    return Response(generate(), mimetype="text/event-stream")
+
+# Search page
+@app.route("/search", methods=["GET"])
+def search_page():
+    storage_mode = request.args.get("storage_mode", "hf")
+    if storage_mode == "hf":
+        try:
+            dataset = load_dataset(HF_DATASET_REPO, token=HF_TOKEN, cache_dir="/app/cache")["train"]
+            files = [{"filename": f, "url": u, "pages": p} for f, u, p in zip(dataset["filename"], dataset["url"], dataset["pages"])]
+            return render_template("search.html", files=files, storage_mode=storage_mode)
+        except Exception as e:
+            logger.error(f"Error loading search page: {e}")
+            return render_template("search.html", files=[], error=str(e), storage_mode=storage_mode)
+    else:  # local
+        files = []
+        results = chroma_collection.get()
+        for i, metadata in enumerate(results["metadatas"]):
+            files.append({
+                "filename": metadata["filename"],
+                "url": "",
+                "pages": [{"page_num": metadata["page_num"], "text": results["documents"][i], "page_file": metadata["page_file"], "words": json.loads(metadata["words"])}]
+            })
+        return render_template("search.html", files=files, storage_mode=storage_mode)
+
+# Semantic search route
+@app.route("/search_documents", methods=["POST"])
+def search_documents():
+    query = request.form.get("query")
+    storage_mode = request.form.get("storage_mode", "hf")
+    if not query:
+        return jsonify({"error": "No query provided"}), 400
+    
+    if storage_mode == "hf":
+        try:
+            dataset = load_dataset(HF_DATASET_REPO, token=HF_TOKEN, cache_dir="/app/cache")["train"]
+            query_embedding = embedder.encode(query).tolist()
+            
+            embeddings = [e for e in dataset["embedding"] if e is not None]
+            documents = dataset["pages"]
+            filenames = dataset["filename"]
+            urls = dataset["url"]
+            processed_ats = dataset["processed_at"]
+            pdf_page_offsets = dataset["pdf_page_offset"]
+            
+            similarities = [
+                dot(query_embedding, emb) / (norm(query_embedding) * norm(emb)) if norm(emb) != 0 else 0
+                for emb in embeddings
+            ]
+            
+            sorted_indices = sorted(range(len(similarities)), key=lambda i: similarities[i], reverse=True)[:5]
+            results = []
+            
+            for idx, i in enumerate(sorted_indices):
+                pages = documents[i]
+                highlighted_pages = []
+                for page in pages:
+                    words = page["words"]
+                    text = page["text"]
+                    pdf_page_num = page["pdf_page"]
+                    page_file = page["page_file"]
+                    page_url = f"https://huggingface.co/datasets/{HF_DATASET_REPO}/resolve/main/{page_file}"
+                    response = requests.get(page_url)
+                    response.raise_for_status()
+                    pdf_bytes = response.content
+                    pdf_base64 = base64.b64encode(pdf_bytes).decode('utf-8')
+                    
+                    sentences = re.split(r'(?<=[.!?])\s+', text)
+                    highlights = []
+                    for sent_idx, sentence in enumerate(sentences):
+                        sent_embedding = embedder.encode(sentence).tolist()
+                        similarity = dot(query_embedding, sent_embedding) / (norm(query_embedding) * norm(sent_embedding)) if norm(sent_embedding) != 0 else 0
+                        if similarity > 0.7:
+                            matching_words = []
+                            sent_words = sentence.split()
+                            word_idx = 0
+                            for word in words:
+                                if word_idx < len(sent_words) and word["text"].lower() in sent_words[word_idx].lower():
+                                    matching_words.append(word)
+                                    word_idx += 1
+                            highlights.append({"sentence": sentence, "index": sent_idx, "words": matching_words})
+                    highlighted_pages.append({
+                        "page_num": page["page_num"],
+                        "text": text,
+                        "highlights": highlights,
+                        "pdf_page": pdf_page_num,
+                        "pdf_data": pdf_base64,
+                        "page_url": page_url
+                    })
+                results.append({
+                    "filename": filenames[i],
+                    "pages": highlighted_pages,
+                    "url": urls[i],
+                    "processed_at": processed_ats[i],
+                    "similarity": similarities[i],
+                    "pdf_page_offset": pdf_page_offsets[i]
+                })
+            return jsonify({"results": results})
+        except Exception as e:
+            logger.error(f"Search error: {e}")
+            return jsonify({"error": str(e)}), 500
+    else:  # local with ChromaDB
+        try:
+            query_results = chroma_collection.query(query_texts=[query], n_results=5)
+            results = []
+            for i, doc in enumerate(query_results["documents"][0]):
+                metadata = query_results["metadatas"][0][i]
+                words = json.loads(metadata["words"])
+                text = doc
+                sentences = re.split(r'(?<=[.!?])\s+', text)
+                highlights = []
+                query_embedding = embedder.encode(query).tolist()
+                for sent_idx, sentence in enumerate(sentences):
+                    sent_embedding = embedder.encode(sentence).tolist()
+                    similarity = dot(query_embedding, sent_embedding) / (norm(query_embedding) * norm(sent_embedding)) if norm(sent_embedding) != 0 else 0
+                    if similarity > 0.7:
+                        matching_words = []
+                        sent_words = sentence.split()
+                        word_idx = 0
+                        for word in words:
+                            if word_idx < len(sent_words) and word["text"].lower() in sent_words[word_idx].lower():
+                                matching_words.append(word)
+                                word_idx += 1
+                        highlights.append({"sentence": sentence, "index": sent_idx, "words": matching_words})
+                with open(metadata["page_file"], "rb") as f:
+                    pdf_bytes = f.read()
+                pdf_base64 = base64.b64encode(pdf_bytes).decode('utf-8')
+                results.append({
+                    "filename": metadata["filename"],
+                    "pages": [{
+                        "page_num": metadata["page_num"],
+                        "text": text,
+                        "highlights": highlights,
+                        "pdf_page": metadata["page_num"],
+                        "pdf_data": pdf_base64,
+                        "page_url": metadata["page_file"]
+                    }],
+                    "url": "",
+                    "processed_at": datetime.now().isoformat(),
+                    "similarity": query_results["distances"][0][i]
+                })
+            return jsonify({"results": results})
+        except Exception as e:
+            logger.error(f"ChromaDB search error: {e}")
+            return jsonify({"error": str(e)}), 500
+
+# Download output folder
+@app.route("/download_output", methods=["GET"])
+def download_output():
+    try:
+        zip_path = "/app/output.zip"
+        shutil.make_archive("/app/output", "zip", OUTPUT_DIR)
+        return send_file(zip_path, download_name="output.zip", as_attachment=True, mimetype="application/zip")
+    except Exception as e:
+        logger.error(f"Error creating zip: {e}")
+        return jsonify({"error": str(e)}), 500
+
+# Preview output contents
+@app.route("/preview_output", methods=["GET"])
+def preview_output():
+    try:
+        combined_pdf_base64 = ""
+        if os.path.exists(COMBINED_PDF_PATH):
+            with open(COMBINED_PDF_PATH, "rb") as f:
+                combined_pdf_base64 = base64.b64encode(f.read()).decode('utf-8')
+        
+        progress_json = {}
+        if os.path.exists(PROGRESS_JSON_PATH):
+            with open(PROGRESS_JSON_PATH, "r") as f:
+                progress_json = json.load(f)
+        
+        return jsonify({
+            "combined_pdf": combined_pdf_base64,
+            "progress_json": progress_json
+        })
     except Exception as e:
-        return jsonify({"status": "error", "message": f"Error loading settings: {str(e)}"}), 500
+        logger.error(f"Error previewing output: {e}")
+        return jsonify({"error": str(e)}), 500
 
-if __name__ == '__main__':
-    port = int(os.environ.get('PORT', 7860))  # Default port for Hugging Face Spaces
-    app.run(host='0.0.0.0', port=port)
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))
+    app.run(host="0.0.0.0", port=port, debug=True)