CoolFace
Apppublic

codeNJ/Instinctive-Studio-assessment

sourceHugging Faceupdated 1y agoView on Hugging Face
0likes
app.py956 linesDownload Raw Back to root
1from fastapi import FastAPI, HTTPException2from pydantic import BaseModel3from typing import List, Optional, Dict, Any4import zipfile5import json6import os7import sqlite38import PyPDF29import re10from sentence_transformers import SentenceTransformer11import faiss12import numpy as np13from sklearn.feature_extraction.text import TfidfVectorizer14from sklearn.metrics.pairwise import cosine_similarity15import nltk16from nltk.tokenize import sent_tokenize17import time18import traceback19import uvicorn20import gradio as gr21from fastapi.middleware.cors import CORSMiddleware22from fastapi.responses import HTMLResponse, JSONResponse23import asyncio24from contextlib import asynccontextmanager25import requests26import urllib.parse27from pathlib import Path28 29# Download NLTK data for sentence tokenization30nltk.download('punkt', quiet=True)31 32# Load sources.json if it exists33try:34    with open('sources.json', 'r') as f:35        sources = json.load(f)36except FileNotFoundError:37    print("sources.json not found, using empty sources list")38    sources = []39 40# Pydantic models for request/response41class AskRequest(BaseModel):42    q: str43    k: int = 544    mode: str = "answer"  # "answer" or "search"45 46class Context(BaseModel):47    text: str48    score: float49    title: str50    url: str51    vector_score: float52    keyword_score: float53    title_score: float54    hybrid_score: float55    cosine_score: float56    exact_matches: int57 58class AskResponse(BaseModel):59    answer: Optional[str] = None60    contexts: List[Context]61    reranker_used: bool = True62 63# Global system instance64safety_system = None65 66@asynccontextmanager67async def lifespan(app: FastAPI):68    # Startup69    global safety_system70    print("Setting up industrial safety retrieval system...")71    safety_system = IndustrialSafetyRetrievalSystem()72    print("Downloading and processing documents...")73    74    # Run processing in a separate thread to avoid blocking75    import threading76    def process_docs():77        safety_system.download_and_process_documents()78        print("System ready!")79    80    thread = threading.Thread(target=process_docs)81    thread.daemon = True82    thread.start()83    84    yield  # App runs here85    86    # Shutdown87    # Clean up resources if needed88    print("Shutting down...")89 90app = FastAPI(91    title="Industrial Safety Retrieval API",92    description="API for retrieving industrial safety information from PDF documents",93    lifespan=lifespan94)95 96# Add CORS middleware97app.add_middleware(98    CORSMiddleware,99    allow_origins=["*"],100    allow_credentials=True,101    allow_methods=["*"],102    allow_headers=["*"],103)104 105class IndustrialSafetyRetrievalSystem:106    def __init__(self):107        print("Loading sentence transformer model...")108        self.model = SentenceTransformer('all-MiniLM-L6-v2')109        print("Model loaded successfully!")110        self.vectorizer = TfidfVectorizer(stop_words='english')111        self.db_path = 'safety_docs.db'112        self.chunk_id_mapping = []  # Maps FAISS index to database chunk ID113        self.setup_database()114 115    def get_db_connection(self):116        """Get a new database connection with proper settings"""117        conn = sqlite3.connect(self.db_path)118        conn.execute("PRAGMA busy_timeout = 30000")119        return conn120 121    def setup_database(self):122        """Create database tables for chunks and metadata"""123        print("Setting up database...")124        conn = self.get_db_connection()125        cursor = conn.cursor()126 127        cursor.execute('''128            CREATE TABLE IF NOT EXISTS sources (129                id INTEGER PRIMARY KEY,130                title TEXT,131                url TEXT,132                file_path TEXT133            )134        ''')135 136        cursor.execute('''137            CREATE TABLE IF NOT EXISTS chunks (138                id INTEGER PRIMARY KEY,139                source_id INTEGER,140                chunk_text TEXT,141                page_number INTEGER,142                chunk_index INTEGER,143                FOREIGN KEY (source_id) REFERENCES sources (id)144            )145        ''')146 147        cursor.execute('''148            CREATE TABLE IF NOT EXISTS embeddings (149                chunk_id INTEGER PRIMARY KEY,150                embedding BLOB,151                FOREIGN KEY (chunk_id) REFERENCES chunks (id)152            )153        ''')154 155        conn.commit()156        conn.close()157        print("Database setup complete!")158 159    def clear_database(self):160        """Clear existing data from database"""161        max_retries = 5162        for attempt in range(max_retries):163            try:164                conn = self.get_db_connection()165                cursor = conn.cursor()166 167                # Disable foreign keys temporarily for faster deletion168                cursor.execute("PRAGMA foreign_keys = OFF")169 170                # Clear tables in reverse order to avoid foreign key constraints171                cursor.execute("DELETE FROM embeddings")172                cursor.execute("DELETE FROM chunks")173                cursor.execute("DELETE FROM sources")174 175                # Re-enable foreign keys176                cursor.execute("PRAGMA foreign_keys = ON")177 178                conn.commit()179                conn.close()180                print("Database cleared successfully")181 182                # Clear the mapping as well183                self.chunk_id_mapping = []184                return True185 186            except sqlite3.OperationalError as e:187                print(f"Database locked (attempt {attempt + 1}/{max_retries}), retrying...")188                time.sleep(2)  # Wait before retrying189                if attempt == max_retries - 1:190                    print(f"Failed to clear database after {max_retries} attempts: {e}")191                    return False192 193    def download_pdfs_from_sources(self):194        """Download PDFs from sources.json and save to pdf_docs directory"""195        pdf_dir = 'pdf_docs'196        os.makedirs(pdf_dir, exist_ok=True)197        198        downloaded_files = []199        200        for source in sources:201            url = source.get('url', '')202            title = source.get('title', 'Unknown')203            204            if not url.lower().endswith('.pdf'):205                print(f"Skipping non-PDF source: {title} ({url})")206                continue207                208            try:209                # Create a safe filename from the title210                safe_title = "".join(c for c in title if c.isalnum() or c in (' ', '-', '_')).rstrip()211                safe_title = safe_title.replace(' ', '_')[:100]  # Limit length212                filename = f"{safe_title}.pdf"213                file_path = os.path.join(pdf_dir, filename)214                215                # Check if file already exists216                if os.path.exists(file_path):217                    print(f"File already exists: {filename}")218                    downloaded_files.append({219                        'file_path': file_path,220                        'file_name': filename,221                        'title': title,222                        'url': url223                    })224                    continue225                226                print(f"Downloading: {title} from {url}")227                228                # Download the PDF229                response = requests.get(url, timeout=30)230                response.raise_for_status()231                232                # Save the PDF233                with open(file_path, 'wb') as f:234                    f.write(response.content)235                236                print(f"Downloaded: {filename} ({len(response.content)} bytes)")237                238                downloaded_files.append({239                    'file_path': file_path,240                    'file_name': filename,241                    'title': title,242                    'url': url243                })244                245            except Exception as e:246                print(f"Error downloading {title} from {url}: {e}")247        248        return downloaded_files249 250    def extract_text_from_pdf(self, file_path):251        """Extract text from PDF file"""252        text = ""253        try:254            print(f"  Opening PDF: {os.path.basename(file_path)}")255            with open(file_path, 'rb') as file:256                pdf_reader = PyPDF2.PdfReader(file)257 258                if pdf_reader.is_encrypted:259                    print("  PDF is encrypted, trying to decrypt...")260                    try:261                        pdf_reader.decrypt('')262                        print("  PDF decrypted successfully!")263                    except:264                        print("  Could not decrypt PDF")265                        return text266 267                total_pages = len(pdf_reader.pages)268                print(f"  Total pages: {total_pages}")269 270                successful_pages = 0271                for page_num, page in enumerate(pdf_reader.pages):272                    try:273                        page_text = page.extract_text()274                        if page_text and page_text.strip():275                            text += f"\n--- PAGE {page_num + 1} ---\n"276                            text += page_text277                            successful_pages += 1278                        else:279                            text += f"\n--- PAGE {page_num + 1} (No text content) ---\n"280                    except Exception as page_error:281                        print(f"  Error on page {page_num + 1}: {page_error}")282                        text += f"\n--- PAGE {page_num + 1} (Error reading page) ---\n"283 284                print(f"  Successfully extracted text from {successful_pages}/{total_pages} pages")285                print(f"  Total text length: {len(text)} characters")286 287        except Exception as e:288            print(f"  Error reading PDF: {e}")289            print(f"  Traceback: {traceback.format_exc()}")290            return ""291 292        return text293 294    def chunk_text(self, text, max_chunk_size=500):295        """Split text into sensible chunks"""296        chunks = []297 298        if not text or not text.strip():299            print("  No text to chunk")300            return chunks301 302        # Clean the text303        text = re.sub(r'--- PAGE \d+ ---', '', text)304        text = re.sub(r'\(No text content\)', '', text)305        text = re.sub(r'\(Error reading page\)', '', text)306        text = text.strip()307 308        if not text:309            print("  Text is empty after cleaning")310            return chunks311 312        paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()]313        print(f"  Found {len(paragraphs)} paragraphs")314 315        for paragraph in paragraphs:316            if len(paragraph) <= max_chunk_size:317                chunks.append(paragraph)318            else:319                try:320                    sentences = sent_tokenize(paragraph)321                    current_chunk = ""322                    for sentence in sentences:323                        if len(current_chunk) + len(sentence) <= max_chunk_size:324                            current_chunk += " " + sentence if current_chunk else sentence325                        else:326                            if current_chunk:327                                chunks.append(current_chunk)328                            current_chunk = sentence329                    if current_chunk:330                        chunks.append(current_chunk)331                except:332                    # Fallback: split by words333                    words = paragraph.split()334                    current_chunk = []335                    current_length = 0336 337                    for word in words:338                        if current_length + len(word) + 1 <= max_chunk_size:339                            current_chunk.append(word)340                            current_length += len(word) + 1341                        else:342                            if current_chunk:343                                chunks.append(" ".join(current_chunk))344                            current_chunk = [word]345                            current_length = len(word)346 347                    if current_chunk:348                        chunks.append(" ".join(current_chunk))349 350        print(f"  Created {len(chunks)} chunks")351        if chunks:352            print(f"  Sample chunk: {chunks[0][:100]}...")353        return chunks354 355    def get_pdf_files(self):356        """Get all PDF files from the pdf_docs directory"""357        pdf_files = []358        pdf_dir = 'pdf_docs'359 360        if not os.path.exists(pdf_dir):361            print(f"Error: Directory '{pdf_dir}' not found!")362            return pdf_files363 364        for file_name in os.listdir(pdf_dir):365            if file_name.endswith('.pdf'):366                file_path = os.path.join(pdf_dir, file_name)367                368                # Get source info from sources.json369                title, url = self.get_source_info(file_name)370                371                pdf_files.append({372                    'file_path': file_path,373                    'file_name': file_name,374                    'title': title,375                    'url': url376                })377 378        return pdf_files379 380    def debug_database_state(self):381        """Debug function to check database state"""382        conn = self.get_db_connection()383        cursor = conn.cursor()384 385        print("\n=== DATABASE DEBUG INFO ===")386 387        # Check sources388        cursor.execute("SELECT COUNT(*) FROM sources")389        source_count = cursor.fetchone()[0]390        print(f"Sources in database: {source_count}")391 392        # Check chunks393        cursor.execute("SELECT COUNT(*) FROM chunks")394        chunk_count = cursor.fetchone()[0]395        print(f"Chunks in database: {chunk_count}")396 397        # Check embeddings398        cursor.execute("SELECT COUNT(*) FROM embeddings")399        embedding_count = cursor.fetchone()[0]400        print(f"Embeddings in database: {embedding_count}")401 402        # Check mapping403        print(f"Chunk ID mapping length: {len(self.chunk_id_mapping)}")404 405        # Show sample data406        if chunk_count > 0:407            cursor.execute("SELECT chunk_text FROM chunks LIMIT 1")408            sample_chunk = cursor.fetchone()409            if sample_chunk:410                print(f"Sample chunk text: {sample_chunk[0][:100]}...")411 412        conn.close()413        print("===========================\n")414 415    def get_source_info(self, file_name):416        """Get proper source info from sources.json"""417        # Try to match the filename with sources.json entries418        file_key = file_name.replace('.pdf', '').replace('_', ' ').replace('-', ' ').lower()419 420        for source in sources:421            source_title = source.get('title', '').lower()422            source_url = source.get('url', 'unknown')423            424            # Check if filename matches title or URL425            if (file_key in source_title or 426                any(word in source_title for word in file_key.split() if len(word) > 3) or427                file_name in source_url):428                return source.get('title', file_name), source.get('url', 'unknown')429 430        # Fallback to filename431        return file_name.replace('.pdf', ''), 'unknown'432 433    def download_and_process_documents(self):434        """Download PDFs from sources.json and process them"""435        print("Starting document download and processing...")436 437        # Clear database first438        if not self.clear_database():439            print("Warning: Could not clear database, proceeding anyway...")440 441        # Download PDFs from sources.json442        pdf_files = self.download_pdfs_from_sources()443        444        # Also check for existing PDFs in the directory445        existing_pdfs = self.get_pdf_files()446        447        # Combine downloaded and existing files448        all_pdf_files = pdf_files + existing_pdfs449        450        # Remove duplicates451        seen_paths = set()452        unique_pdf_files = []453        for pdf in all_pdf_files:454            if pdf['file_path'] not in seen_paths:455                seen_paths.add(pdf['file_path'])456                unique_pdf_files.append(pdf)457 458        if not unique_pdf_files:459            print("No PDF files found to process!")460            return461 462        print(f"Found {len(unique_pdf_files)} PDF files to process")463 464        all_embeddings = []465        processed_files = 0466 467        # Process each PDF file468        for pdf_info in unique_pdf_files:469            file_path = pdf_info['file_path']470            file_name = pdf_info['file_name']471            title = pdf_info.get('title', '')472            url = pdf_info.get('url', 'unknown')473 474            if not os.path.exists(file_path):475                print(f"File not found: {file_path}")476                continue477 478            print(f"\n=== Processing: {file_name} ===")479 480            # Get a new database connection for this file481            conn = self.get_db_connection()482            cursor = conn.cursor()483 484            try:485                # If title wasn't provided, try to get it from sources.json486                if not title:487                    title, url = self.get_source_info(file_name)488 489                # Insert source into database490                cursor.execute(491                    "INSERT INTO sources (title, url, file_path) VALUES (?, ?, ?)",492                    (title, url, file_path)493                )494                source_id = cursor.lastrowid495 496                # Extract and chunk text497                text = self.extract_text_from_pdf(file_path)498                if not text or not text.strip():499                    print(f"  No text extracted from {file_name}")500                    conn.close()501                    continue502 503                chunks = self.chunk_text(text)504                if not chunks:505                    print(f"  No chunks created from {file_name}")506                    conn.close()507                    continue508 509                # Store chunks and generate embeddings510                for chunk_idx, chunk_text in enumerate(chunks):511                    cursor.execute(512                        "INSERT INTO chunks (source_id, chunk_text, page_number, chunk_index) VALUES (?, ?, ?, ?)",513                        (source_id, chunk_text, 1, chunk_idx)514                    )515                    chunk_id = cursor.lastrowid516 517                    # Generate embedding518                    embedding = self.model.encode([chunk_text])[0]519 520                    # Store embedding in database521                    cursor.execute(522                        "INSERT INTO embeddings (chunk_id, embedding) VALUES (?, ?)",523                        (chunk_id, embedding.tobytes())524                    )525 526                    # Add to our collections for FAISS527                    all_embeddings.append(embedding)528                    self.chunk_id_mapping.append(chunk_id)  # Map FAISS index to chunk ID529 530                conn.commit()531                processed_files += 1532                print(f"  Successfully processed {file_name}")533 534            except Exception as e:535                print(f"  Error processing {file_name}:")536                print(f"    Error: {e}")537                print(f"    Traceback: {traceback.format_exc()}")538                conn.rollback()539            finally:540                conn.close()541 542        # Debug database state after processing543        self.debug_database_state()544 545        # Create FAISS index with all embeddings546        if all_embeddings:547            try:548                print(f"\nCreating FAISS index with {len(all_embeddings)} embeddings...")549                embeddings_array = np.array(all_embeddings).astype('float32')550                print(f"Embeddings array shape: {embeddings_array.shape}")551                print(f"Chunk ID mapping length: {len(self.chunk_id_mapping)}")552 553                # Normalize embeddings for better cosine similarity554                faiss.normalize_L2(embeddings_array)555 556                # Use Inner Product index for normalized vectors (equivalent to cosine similarity)557                self.faiss_index = faiss.IndexFlatIP(embeddings_array.shape[1])558                self.faiss_index.add(embeddings_array)559                print(f"FAISS index created with {self.faiss_index.ntotal} vectors (using cosine similarity)")560 561            except Exception as e:562                print(f"Error creating FAISS index: {e}")563                print(f"FAISS traceback: {traceback.format_exc()}")564        else:565            print("No embeddings were generated!")566 567        print(f"\nProcessing complete: {processed_files}/{len(unique_pdf_files)} files with {len(all_embeddings)} chunks")568 569    def search(self, query, top_k=10):570        """Search for relevant chunks with improved scoring"""571        print(f"\nSearching for: '{query}'")572 573        if not hasattr(self, 'faiss_index') or self.faiss_index.ntotal == 0:574            print("No FAISS index available!")575            return []576 577        print(f"FAISS index has {self.faiss_index.ntotal} vectors")578        print(f"Chunk ID mapping has {len(self.chunk_id_mapping)} entries")579 580        # Vector similarity search with normalized query581        query_embedding = self.model.encode([query])[0].astype('float32').reshape(1, -1)582        faiss.normalize_L2(query_embedding)  # Normalize query for cosine similarity583        scores, indices = self.faiss_index.search(query_embedding, min(top_k * 3, self.faiss_index.ntotal))584 585        print(f"FAISS search returned {len(indices[0])} results")586 587        # Get chunk details using the mapping588        conn = self.get_db_connection()589        cursor = conn.cursor()590        results = []591 592        for faiss_idx, score in zip(indices[0], scores[0]):593            # Skip invalid indices594            if faiss_idx >= len(self.chunk_id_mapping) or faiss_idx < 0:595                print(f"Skipping invalid FAISS index: {faiss_idx}")596                continue597 598            # Get the actual chunk ID from our mapping599            chunk_id = self.chunk_id_mapping[faiss_idx]600 601            cursor.execute('''602                SELECT c.id, c.chunk_text, s.title, s.url603                FROM chunks c604                JOIN sources s ON c.source_id = s.id605                WHERE c.id = ?606            ''', (chunk_id,))607 608            row = cursor.fetchone()609            if row:610                chunk_id, chunk_text, title, url = row611 612                # Vector score is now cosine similarity (0 to 1, higher is better)613                vector_score = max(0, float(score))  # Ensure non-negative614 615                # Enhanced keyword scoring616                query_terms = [term.lower().strip() for term in query.split() if len(term.strip()) > 2]617                chunk_text_lower = chunk_text.lower()618 619                # Exact matches620                exact_matches = sum(1 for term in query_terms if term in chunk_text_lower)621 622                # Partial matches (stemming-like)623                partial_matches = 0624                for term in query_terms:625                    if len(term) > 4:  # Only for longer terms626                        stem = term[:4]  # Simple stemming627                        if stem in chunk_text_lower and term not in chunk_text_lower:628                            partial_matches += 0.5629 630                # Calculate keyword score with bonuses631                if query_terms:632                    keyword_score = (exact_matches + partial_matches) / len(query_terms)633                    # Bonus for multiple matches in close proximity634                    if exact_matches > 1:635                        keyword_score *= 1.2636                else:637                    keyword_score = 0638 639                # Context relevance score (title matching)640                title_score = 0641                if title and query_terms:642                    title_lower = title.lower()643                    title_matches = sum(1 for term in query_terms if term in title_lower)644                    title_score = title_matches / len(query_terms)645 646                # Combined hybrid score with better weighting647                # Increase vector weight since embeddings are generally more reliable648                hybrid_score = (0.6 * vector_score +649                                0.3 * keyword_score +650                                0.1 * title_score)651 652                # Boost score for very relevant chunks653                if keyword_score > 0.5:  # High keyword match654                    hybrid_score *= 1.15655                if vector_score > 0.8:  # High semantic similarity656                    hybrid_score *= 1.1657 658                # Ensure score doesn't exceed 1.0659                hybrid_score = min(hybrid_score, 1.0)660 661                results.append({662                    'chunk_id': chunk_id,663                    'chunk_text': chunk_text,664                    'title': title,665                    'url': url,666                    'vector_score': vector_score,667                    'keyword_score': keyword_score,668                    'title_score': title_score,669                    'hybrid_score': hybrid_score,670                    'cosine_score': score,671                    'exact_matches': exact_matches672                })673            else:674                print(f"No chunk found for ID {chunk_id}")675 676        conn.close()677 678        print(f"Found {len(results)} matching chunks")679        680        results.sort(key=lambda x: x['hybrid_score'], reverse=True)681        return results[:top_k]682 683    def answer_question(self, question, top_k=5, confidence_threshold=0.4):684        """Generate answer based on retrieved chunks with abstention logic"""685        print(f"\nAnswering question: '{question}'")686        relevant_chunks = self.search(question, top_k)687 688        if not relevant_chunks:689            return "I couldn't find any relevant information in the industrial safety documents to answer this question. Please try rephrasing your question or asking about a different safety topic."690 691        # Check if we should abstain based on confidence692        best_score = relevant_chunks[0]['hybrid_score']693        avg_score = sum(chunk['hybrid_score'] for chunk in relevant_chunks[:3]) / min(3, len(relevant_chunks))694 695        print(f"Best score: {best_score:.3f}, Average top-3 score: {avg_score:.3f}")696 697        if best_score < confidence_threshold:698            return f"I found some potentially related content (confidence: {best_score:.3f}), but I'm not confident it adequately addresses your question about '{question}'. The retrieved information may not be specific enough. Please try a more specific question or different keywords related to industrial safety."699 700        print(f"Found {len(relevant_chunks)} relevant chunks with good confidence")701 702        # Create grounded answer703        answer_parts = []704        citations = []705        seen_sources = set()706 707        # Use top 3 chunks for answer708        for i, chunk in enumerate(relevant_chunks[:3]):709            # Add chunk content710            chunk_text = chunk['chunk_text'].strip()711            if chunk_text:712                answer_parts.append(f"According to the safety documentation: {chunk_text}")713 714            # Add unique citations715            source_key = f"{chunk['title']}|{chunk['url']}"716            if source_key not in seen_sources:717                citations.append(f"[{len(citations) + 1}] {chunk['title']} - {chunk['url']}")718                seen_sources.add(source_key)719 720        if not answer_parts:721            return "I found potentially relevant documents but couldn't extract clear information to answer your question. Please try rephrasing with more specific safety-related terms."722 723        # Construct final answer724        answer = f"Based on the industrial safety documentation:\n\n"725        answer += "\n\n".join(answer_parts)726 727        if citations:728            answer += f"\n\nSources:\n" + "\n".join(citations)729 730        # Add confidence note for borderline cases731        if best_score < 0.6:732            answer += f"\n\nNote: Answer confidence is moderate ({best_score:.2f}). Please verify with additional sources if this is for critical safety decisions."733 734        return answer735 736@app.post("/ask", response_model=AskResponse)737async def ask_question(request: AskRequest):738    """Endpoint to ask questions about industrial safety"""739    global safety_system740 741    if safety_system is None:742        raise HTTPException(status_code=503, detail="System not initialized yet")743 744    try:745        if request.mode == "search":746            # Return search results only747            results = safety_system.search(request.q, request.k)748            contexts = [749                Context(750                    text=result['chunk_text'],751                    score=result['hybrid_score'],752                    title=result['title'],753                    url=result['url'],754                    vector_score=result['vector_score'],755                    keyword_score=result['keyword_score'],756                    title_score=result['title_score'],757                    hybrid_score=result['hybrid_score'],758                    cosine_score=result['cosine_score'],759                    exact_matches=result['exact_matches']760                )761                for result in results762            ]763            return AskResponse(answer=None, contexts=contexts, reranker_used=True)764 765        else:  # answer mode766            # Generate answer with contexts767            results = safety_system.search(request.q, request.k)768            answer = safety_system.answer_question(request.q, request.k)769 770            contexts = [771                Context(772                    text=result['chunk_text'],773                    score=result['hybrid_score'],774                    title=result['title'],775                    url=result['url'],776                    vector_score=result['vector_score'],777                    keyword_score=result['keyword_score'],778                    title_score=result['title_score'],779                    hybrid_score=result['hybrid_score'],780                    cosine_score=result['cosine_score'],781                    exact_matches=result['exact_matches']782                )783                for result in results784            ]785 786            return AskResponse(answer=answer, contexts=contexts, reranker_used=True)787 788    except Exception as e:789        raise HTTPException(status_code=500, detail=f"Error processing request: {str(e)}")790 791@app.get("/health")792async def health_check():793    """Health check endpoint"""794    global safety_system795    if safety_system is None:796        return {"status": "initializing"}797    798    if not hasattr(safety_system, 'faiss_index'):799        return {"status": "processing_documents"}800 801    return {802        "status": "ready",803        "documents_processed": len(safety_system.chunk_id_mapping) if safety_system else 0,804        "faiss_index_size": safety_system.faiss_index.ntotal if hasattr(safety_system, 'faiss_index') else 0805    }806 807@app.get("/")808async def root():809    """Root endpoint with API information"""810    return {811        "message": "Industrial Safety Retrieval API",812        "endpoints": {813            "POST /ask": "Ask questions about industrial safety",814            "GET /health": "Check system health status",815            "GET /": "This information page"816        },817        "parameters": {818            "q": "Question to ask",819            "k": "Number of results to return (default: 5)",820            "mode": "'answer' for full answer or 'search' for contexts only (default: 'answer')"821        }822    }823 824# Gradio Interface825def create_gradio_interface():826    """Create Gradio interface for the safety retrieval system"""827    828    def gradio_ask_question(question, num_results=5, mode="answer"):829        """Wrapper function for Gradio interface"""830        if safety_system is None:831            return "System not initialized yet. Please wait...", []832        833        if not hasattr(safety_system, 'faiss_index'):834            return "System is still processing documents. Please wait...", []835        836        try:837            if mode == "answer":838                answer = safety_system.answer_question(question, num_results)839                results = safety_system.search(question, num_results)840                contexts = [841                    {842                        "text": result['chunk_text'][:200] + "..." if len(result['chunk_text']) > 200 else result['chunk_text'],843                        "score": round(result['hybrid_score'], 3),844                        "title": result['title'],845                        "url": result['url'],846                        "confidence": f"{result['hybrid_score']:.3f}"847                    }848                    for result in results849                ]850                return answer, contexts851            else:852                results = safety_system.search(question, num_results)853                contexts = [854                    {855                        "text": result['chunk_text'][:200] + "..." if len(result['chunk_text']) > 200 else result['chunk_text'],856                        "score": round(result['hybrid_score'], 3),857                        "title": result['title'],858                        "url": result['url'],859                        "confidence": f"{result['hybrid_score']:.3f}"860                    }861                    for result in results862                ]863                return "Search results (showing top contexts):", contexts864                865        except Exception as e:866            return f"Error: {str(e)}", []867    868    with gr.Blocks(title="Industrial Safety Q&A", theme=gr.themes.Soft()) as demo:869        gr.Markdown("# 🛡️ Industrial Safety Document Q&A")870        gr.Markdown("Ask questions about industrial safety standards and procedures")871        872        with gr.Row():873            with gr.Column(scale=2):874                question = gr.Textbox(875                    label="Question",876                    placeholder="e.g., What are the electrical safety requirements?",877                    lines=2,878                    max_lines=4879                )880                881                with gr.Row():882                    num_results = gr.Slider(883                        minimum=1,884                        maximum=10,885                        value=5,886                        step=1,887                        label="Number of results"888                    )889                    mode = gr.Radio(890                        choices=["answer", "search"],891                        value="answer",892                        label="Response Mode"893                    )894                895                submit_btn = gr.Button("Ask Question", variant="primary")896            897            with gr.Column(scale=3):898                answer_output = gr.Textbox(899                    label="Answer",900                    lines=6,901                    interactive=False902                )903                904                contexts_output = gr.DataFrame(905                    label="Supporting Contexts",906                    headers=["text", "score", "title", "url", "confidence"]907                )908        909        examples = gr.Examples(910            examples=[911                ["What are the electrical safety requirements?", 5, "answer"],912                ["How should lockout tagout procedures be implemented?", 3, "answer"],913                ["What PPE is required for chemical handling?", 4, "search"],914                ["What are the emergency procedures for fire incidents?", 5, "answer"]915            ],916            inputs=[question, num_results, mode],917            outputs=[answer_output, contexts_output],918            fn=gradio_ask_question,919            cache_examples=False,920            label="Example Queries"921        )922        923        submit_btn.click(924            fn=gradio_ask_question,925            inputs=[question, num_results, mode],926            outputs=[answer_output, contexts_output],927            api_name="ask_question"928        )929        930        question.submit(931            fn=gradio_ask_question,932            inputs=[question, num_results, mode],933            outputs=[answer_output, contexts_output]934        )935    936    return demo937 938gradio_app = create_gradio_interface()939app = gr.mount_gradio_app(app, gradio_app, path="/")940 941@app.get("/gradio")942async def gradio_redirect():943    return HTMLResponse("""944    <!DOCTYPE html>945    <html>946        <head>947            <meta http-equiv="refresh" content="0; url=/" />948        </head>949        <body>950            <p>Redirecting to <a href="/">Gradio Interface</a></p>951        </body>952    </html>953    """)954 955if __name__ == "__main__":956    uvicorn.run(app, host="0.0.0.0", port=7860)