Spaces:
Runtime error
Runtime error
| import asyncio | |
| import hashlib | |
| import logging | |
| import re | |
| import time | |
| from concurrent.futures import ThreadPoolExecutor | |
| from functools import lru_cache | |
| from typing import List, Tuple, Optional | |
| import tempfile | |
| import os | |
| import platform | |
| import faiss | |
| import fitz | |
| import google.generativeai as genai | |
| import nltk | |
| import numpy as np | |
| import requests | |
| from fastapi import FastAPI, HTTPException, Depends, Header | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from pydantic import BaseModel | |
| # ------------------- Cross-Platform NLTK Setup ------------------- | |
| def setup_nltk(): | |
| """Setup NLTK data directory based on platform and environment.""" | |
| if os.getenv('NLTK_DATA'): | |
| nltk_data_dir = os.getenv('NLTK_DATA') | |
| else: | |
| system = platform.system().lower() | |
| if system == "windows": | |
| base_dir = os.getenv('LOCALAPPDATA', tempfile.gettempdir()) | |
| nltk_data_dir = os.path.join(base_dir, "nltk_data") | |
| elif system == "darwin": # macOS | |
| home = os.path.expanduser("~") | |
| base_dir = os.path.join(home, "Library", "Caches") | |
| if os.path.exists(base_dir): | |
| nltk_data_dir = os.path.join(base_dir, "nltk_data") | |
| else: | |
| nltk_data_dir = os.path.join(tempfile.gettempdir(), "nltk_data") | |
| else: # Linux and others | |
| nltk_data_dir = os.path.join(tempfile.gettempdir(), "nltk_data") | |
| os.makedirs(nltk_data_dir, exist_ok=True) | |
| try: | |
| nltk.data.find("tokenizers/punkt") | |
| except LookupError: | |
| try: | |
| nltk.download("punkt", download_dir=nltk_data_dir, quiet=True) | |
| except Exception as e: | |
| logging.warning(f"Failed to download punkt to {nltk_data_dir}: {e}") | |
| try: | |
| nltk.download("punkt", quiet=True) | |
| except Exception: | |
| logging.error("Failed to download punkt data. Tokenization may fail.") | |
| if nltk_data_dir not in nltk.data.path: | |
| nltk.data.path.append(nltk_data_dir) | |
| return nltk_data_dir | |
| # Setup NLTK | |
| nltk_data_dir = setup_nltk() | |
| try: | |
| from nltk.tokenize import sent_tokenize | |
| except ImportError: | |
| logging.error("NLTK tokenizer not available. Install NLTK and download punkt data.") | |
| def sent_tokenize(text): | |
| # fallback if nltk tokenizer unavailable | |
| import re | |
| sentences = re.split(r'[.!?]+', text) | |
| return [s.strip() for s in sentences if s.strip()] | |
| # -------------- SPEED-OPTIMIZED Configuration -------------- | |
| GEMINI_API_KEY = "AIzaSyCYVp6yfDYGfUTGdbeuibEp5X5-WYfRDmg" | |
| SECRET_TOKEN = "bc7aba495dfdaa88e718dec0d0fba29a2d45eaecac3268453778d027c7419081" | |
| MAX_WORKERS = min(64, (os.cpu_count() or 4) * 4) # Max threads for I/O | |
| EMBED_BATCH = 400 # Batch size for embedding calls | |
| CHUNK_WORDS = 600 # Chunk size in words | |
| OVERLAP_WORDS = 30 # Small overlap to catch boundary context, test effect (was 0) | |
| RETRIEVE_K = 4 # Increased from 3 to 4 for more context chunks | |
| SIMILARITY_THRESHOLD = 0.2 # Increased threshold (from 0.15) to reduce noise | |
| MAX_OUTPUT_TOKENS = 120 # Increased from 90 to 120 for richer answers | |
| CONCURRENCY_LIMIT = min(12, MAX_WORKERS // 2) # Tune per resource availability | |
| genai.configure(api_key=GEMINI_API_KEY) | |
| model = genai.GenerativeModel("gemini-1.5-flash") | |
| executor = ThreadPoolExecutor(max_workers=MAX_WORKERS) | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", | |
| ) | |
| logger = logging.getLogger("hackrx") | |
| logger.info(f"Platform: {platform.system()} {platform.release()}") | |
| logger.info(f"NLTK data directory: {nltk_data_dir}") | |
| logger.info(f"Max workers: {MAX_WORKERS}, Concurrency limit: {CONCURRENCY_LIMIT}") | |
| # ---------- SIMPLE FAST CACHE ---------- | |
| class FastCache: | |
| def __init__(self, max_size=1000): | |
| self.store = {} | |
| self.max_size = max_size | |
| def get(self, key: str): | |
| return self.store.get(key) | |
| def set(self, key: str, value): | |
| if len(self.store) >= self.max_size: | |
| # Remove oldest (simple FIFO) | |
| oldest_key = next(iter(self.store)) | |
| self.store.pop(oldest_key) | |
| self.store[key] = value | |
| def clear(self): | |
| self.store.clear() | |
| # Global caches | |
| doc_cache = FastCache(50) | |
| emb_cache = FastCache(3000) | |
| ans_cache = FastCache(1000) | |
| # ---------- FastAPI Setup ---------- | |
| app = FastAPI(title="HackRx LLM API Final", version="1.0.0") | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| allow_credentials=True, | |
| ) | |
| class Req(BaseModel): | |
| documents: str | |
| questions: List[str] | |
| class Resp(BaseModel): | |
| answers: List[str] | |
| async def verify_token(authorization: str = Header()): | |
| if not authorization or not authorization.startswith("Bearer "): | |
| raise HTTPException(401, "Invalid authorization header") | |
| token = authorization.split()[1] | |
| if token != SECRET_TOKEN: | |
| raise HTTPException(401, "Invalid token") | |
| # ---------- ULTRA-FAST PDF Extraction ---------- | |
| async def extract_text(url: str) -> str: | |
| key = hashlib.md5(url.encode()).hexdigest() | |
| cached = doc_cache.get(f"txt_{key}") | |
| if cached: | |
| return cached | |
| def _extract(): | |
| try: | |
| r = requests.get(url, timeout=10) | |
| r.raise_for_status() | |
| with fitz.open(stream=r.content, filetype="pdf") as doc: | |
| text = " ".join(page.get_text() for page in doc if page.get_text().strip()) | |
| return re.sub(r'\s+', ' ', text).strip() | |
| except Exception as e: | |
| logger.error(f"PDF extraction failed: {e}") | |
| raise HTTPException(500, f"Failed to extract PDF: {str(e)}") | |
| text = await asyncio.get_event_loop().run_in_executor(executor, _extract) | |
| if not text: | |
| raise HTTPException(400, "Extracted PDF text is empty") | |
| doc_cache.set(f"txt_{key}", text) | |
| return text | |
| # ---------- FAST Chunking (with small overlap) ---------- | |
| def chunk_text(text: str) -> List[str]: | |
| words = text.split() | |
| chunks = [] | |
| step = CHUNK_WORDS - OVERLAP_WORDS | |
| if step <= 0: | |
| step = CHUNK_WORDS # fallback to no overlap if misconfigured | |
| for i in range(0, len(words), step): | |
| chunk = " ".join(words[i:i + CHUNK_WORDS]) | |
| if chunk.strip(): | |
| chunks.append(chunk) | |
| return chunks | |
| # ---------- ULTRA-FAST Embeddings ---------- | |
| # Increased cache size from 8192 | |
| def embed_one(text: str, task: str): | |
| key = hashlib.md5((task + text).encode()).hexdigest()[:16] | |
| cached = emb_cache.get(key) | |
| if cached: | |
| return cached | |
| try: | |
| emb = genai.embed_content(model="models/embedding-001", content=text, task_type=task)["embedding"] | |
| except Exception as e: | |
| logger.error(f"Embedding error: {e}") | |
| emb = [0.0] * 768 | |
| emb_cache.set(key, emb) | |
| return emb | |
| async def embed_batch(texts: List[str], task: str) -> List[List[float]]: | |
| if not texts: | |
| return [] | |
| sem = asyncio.Semaphore(20) # High concurrency | |
| async def embed_single_async(text): | |
| async with sem: | |
| return await asyncio.get_event_loop().run_in_executor(executor, embed_one, text, task) | |
| return await asyncio.gather(*[embed_single_async(t) for t in texts]) | |
| # ---------- FAST FAISS Index ---------- | |
| def build_index(embeddings: List[List[float]]) -> Tuple[faiss.Index, np.ndarray]: | |
| if not embeddings: | |
| raise ValueError("No embeddings provided") | |
| mat = np.array(embeddings, dtype=np.float32) | |
| faiss.normalize_L2(mat) | |
| idx = faiss.IndexFlatIP(mat.shape[1]) | |
| idx.add(mat) | |
| return idx, mat | |
| # ---------- SOFT KEYWORD BOOSTING (OPTIONAL) ---------- | |
| def soft_keyword_boost(question: str, candidates: List[Tuple[float, str]]) -> List[Tuple[float, str]]: | |
| """ | |
| Given a list of (similarity, chunk) tuples, | |
| boost similarity by a small factor if chunk contains question keywords. | |
| This is soft boosting, not strict filtering. | |
| """ | |
| keywords = set(re.findall(r"\w+", question.lower())) | |
| boosted = [] | |
| for sim, chunk in candidates: | |
| chunk_words = set(re.findall(r"\w+", chunk.lower())) | |
| common = keywords.intersection(chunk_words) | |
| boost_factor = 1.0 + 0.05 * len(common) # 5% boost per keyword overlap, capped elsewhere | |
| boosted.append((sim * boost_factor, chunk)) | |
| # Sort descending by boosted similarity score | |
| boosted.sort(key=lambda x: x[0], reverse=True) | |
| return boosted | |
| # ---------- FAST Retrieval with soft keyword boosting ---------- | |
| async def retrieve(query: str, chunks: List[str], index: faiss.Index) -> List[str]: | |
| if not query or not chunks: | |
| return [] | |
| q_emb = (await embed_batch([query], "retrieval_query"))[0] | |
| q_vec = np.array([q_emb], dtype=np.float32) | |
| faiss.normalize_L2(q_vec) | |
| top_k = min(RETRIEVE_K * 3, len(chunks)) # Search more initially for boosting | |
| D, I = index.search(q_vec, top_k) | |
| candidates = [] | |
| for score, idx in zip(D[0], I[0]): | |
| if score > SIMILARITY_THRESHOLD: | |
| candidates.append((score, chunks[idx])) | |
| # Apply optional soft keyword boosting for ranking | |
| boosted = soft_keyword_boost(query, candidates) | |
| # Select top RETRIEVE_K chunks after boosting | |
| selected = [chunk for _, chunk in boosted[:RETRIEVE_K]] | |
| return selected | |
| # ---------- FAST Answer Generation ---------- | |
| async def generate(question: str, chunks: List[str]) -> str: | |
| if not chunks: | |
| return "Information not found in the policy." | |
| cache_key = hashlib.md5((question + "".join(chunks)).encode()).hexdigest()[:16] | |
| cached = ans_cache.get(cache_key) | |
| if cached: | |
| return cached | |
| # Compose prompt with selected chunks | |
| context = "\n\n".join(chunks) | |
| prompt = f"""Context: {context} | |
| Question: {question} | |
| Answer briefly and precisely. If not found, say "Information not found in the policy." | |
| Answer:""" | |
| def _gen(): | |
| try: | |
| resp = model.generate_content( | |
| prompt, | |
| generation_config=genai.types.GenerationConfig( | |
| temperature=0.1, | |
| max_output_tokens=MAX_OUTPUT_TOKENS, | |
| top_p=0.7, # Lower for accuracy | |
| top_k=20, | |
| ), | |
| ) | |
| return re.sub(r'\s+', ' ', resp.text.strip()) or "Information not found in the policy." | |
| except Exception as e: | |
| logger.error(f"Generation error: {e}") | |
| return "Information not found in the policy." | |
| answer = await asyncio.get_event_loop().run_in_executor(executor, _gen) | |
| ans_cache.set(cache_key, answer) | |
| return answer | |
| # ---------- MAIN PIPELINE ---------- | |
| async def process(doc_url: str, questions: List[str]) -> List[str]: | |
| start = time.time() | |
| text = await extract_text(doc_url) | |
| chunks = chunk_text(text) | |
| logger.info(f"Extracted and chunked in {time.time() - start:.2f}s ({len(chunks)} chunks)") | |
| idx_key = hashlib.md5("".join(chunks).encode()).hexdigest()[:16] | |
| cached_idx = doc_cache.get(f"idx_{idx_key}") | |
| if cached_idx: | |
| index, mat = cached_idx | |
| logger.info("Using cached index") | |
| else: | |
| emb_start = time.time() | |
| embeddings = await embed_batch(chunks, "retrieval_document") | |
| logger.info(f"Embeddings took {time.time() - emb_start:.2f}s") | |
| index, mat = build_index(embeddings) | |
| doc_cache.set(f"idx_{idx_key}", (index, mat)) | |
| sem = asyncio.Semaphore(CONCURRENCY_LIMIT) | |
| async def handle_q(q): | |
| async with sem: | |
| ctx = await retrieve(q, chunks, index) | |
| return await generate(q, ctx) | |
| answers = await asyncio.gather(*[handle_q(q) for q in questions]) | |
| total_time = time.time() - start | |
| logger.info(f"Total processing time: {total_time:.2f}s for {len(questions)} questions") | |
| return answers | |
| # ---------- Endpoints ---------- | |
| async def run(req: Req, token=Depends(verify_token)): | |
| return {"answers": await process(req.documents, req.questions)} | |
| async def health(): | |
| return { | |
| "status": "ok", | |
| "platform": platform.system(), | |
| "python_version": platform.python_version(), | |
| "nltk_data_dir": nltk_data_dir, | |
| "cache_stats": { | |
| "doc_cache": len(doc_cache.store), | |
| "emb_cache": len(emb_cache.store), | |
| "ans_cache": len(ans_cache.store), | |
| } | |
| } | |
| async def root(): | |
| return {"message": "HackRx LLM API Final is running"} | |
| if __name__ == "__main__": | |
| import uvicorn | |
| uvicorn.run(app, host="0.0.0.0", port=8000, workers=1) | |