diff --git "a/page_files/categorized/Backend/PDF_DataExtraction.py" "b/page_files/categorized/Backend/PDF_DataExtraction.py" --- "a/page_files/categorized/Backend/PDF_DataExtraction.py" +++ "b/page_files/categorized/Backend/PDF_DataExtraction.py" @@ -1,2427 +1,316 @@ """ -DocToDB — Materials Science PDF Extractor (Dual LLM Consensus Pipeline) -======================================================================= -PRIMARY: Gemini 2.0 Flash (first preference, no 1.5) -SECONDARY: GPT-4o - -FLOW: - 1. Extract all chunks (tables + text) - 2. Index into ChromaDB for semantic ranking - 3. Rank ALL chunks via single schema-derived retrieval query - 4. Build batches - 5. Run Gemini + GPT-4o in PARALLEL (ThreadPoolExecutor) - - Each LLM respects its own internal rate-limit delays independently - - No shared state → no interference between the two APIs - 6. Consensus filter — only rows matching in BOTH results are kept - (fuzzy: same property_name+section+material_name, value within 5%) - 7. DOI auto-extracted from PDF by both LLMs; merged with override support - 8. DOI rendered as clickable link in Streamlit UI - -Rate limit handling: - - Gemini: exponential backoff 60s→120s→180s, 65s between batches - - GPT-4o: exponential backoff 30s→60s→90s, 10s between batches - - ChromaDB persistence (no re-embedding for same PDF) - - JSON cache (no re-calling LLMs for same PDF) +FinalVerdict.py +======================= +Straight one-call Gemini extraction (no chunking, no ranking, no batching) with +source verification: every extracted property's value must be found in the PDF +first (tolerant numeric match), then its property+material meaning is checked +against the sentence(s) that contained that value. + +Output: one table — Gemini's extracted properties, each flagged source_verified +True/False, downloadable as CSV. + +Run: streamlit run FinalVerdict.py +Needs GEMINI_API_KEY in .env """ -from __future__ import annotations - -import hashlib -import io -import json -import logging import os import re -import tempfile -import time -from concurrent.futures import ThreadPoolExecutor, as_completed -from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Tuple +import json +import math +import base64 +from typing import Any, Dict, List, Optional -import fitz # pymupdf +import fitz # PyMuPDF import numpy as np import pandas as pd -import pdfplumber import requests +import streamlit as st from dotenv import load_dotenv -import threading - -load_dotenv() # loads .env from the working directory - - -# ── optional deps ───────────────────────────────────────────────────────────── -try: - from docling.document_converter import DocumentConverter - DOCLING_AVAILABLE = True -except ImportError: - DOCLING_AVAILABLE = False - -try: - import chromadb - CHROMA_AVAILABLE = True -except ImportError: - CHROMA_AVAILABLE = False - logging.warning("chromadb not installed — pip install chromadb") - -try: - from sentence_transformers import SentenceTransformer - ST_AVAILABLE = True -except ImportError: - ST_AVAILABLE = False - logging.warning("sentence-transformers not installed — pip install sentence-transformers") - -try: - import camelot - CAMELOT_AVAILABLE = True -except ImportError: - CAMELOT_AVAILABLE = False - -try: - import pytesseract - from PIL import Image - OCR_AVAILABLE = True -except ImportError: - OCR_AVAILABLE = False +load_dotenv() -# ───────────────────────────────────────────────────────────────────────────── -# CONFIG -# ───────────────────────────────────────────────────────────────────────────── - -logging.basicConfig(level=logging.INFO, format="%(levelname)s │ %(message)s") -log = logging.getLogger(__name__) - -if OCR_AVAILABLE: - pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe" - -# ── API Keys ────────────────────────────────────────────────────────────────── GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "") -OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "").strip() - -# ── Gemini model selection (no 1.5 models) ──────────────────────────────────── -_cached_gemini_model: Optional[str] = None - -GEMINI_PREFERRED_MODELS = [ - "gemini-3.5-flash", # GA, 1M context, 65k output — first choice - "gemini-3.1-flash-lite", # fast/cheap fallback - "gemini-3.1-pro-preview", # pro fallback - "gemini-2.5-flash", # last resort only - # gemini-2.0-*, gemini-1.5-* excluded — deprecated/gone -] - -def get_gemini_model(api_key: str) -> str: - global _cached_gemini_model - if _cached_gemini_model: - return _cached_gemini_model - try: - resp = requests.get( - f"https://generativelanguage.googleapis.com/v1beta/models?key={api_key}", - timeout=10, - ) - if not resp.ok: - _cached_gemini_model = "gemini-2.0-flash" - _run_startup_probes(_cached_gemini_model) - return _cached_gemini_model - models = resp.json().get("models", []) - available = [ - m["name"].replace("models/", "") - for m in models - if "generateContent" in m.get("supportedGenerationMethods", []) - ] - for pref in GEMINI_PREFERRED_MODELS: - for name in available: - if name.startswith(pref): - log.info(f"Gemini auto-selected: {name}") - _cached_gemini_model = name - return _cached_gemini_model - flash_models = [ - n for n in available - if "flash" in n - and "1.5" not in n - and "2.0" not in n - and "tts" not in n # exclude audio models - and "audio" not in n # exclude audio models - and "image" not in n # exclude image-gen models - ] - _cached_gemini_model = flash_models[0] if flash_models else "gemini-3.2-flash" - _run_startup_probes(_cached_gemini_model) - return _cached_gemini_model - except Exception as e: - log.warning(f"Gemini model auto-detect failed: {e}") - _cached_gemini_model = "gemini-2.0-flash" - _run_startup_probes(_cached_gemini_model) - return _cached_gemini_model - -GEMINI_MODEL = None -GEMINI_API_URL = None - -# ── OpenAI / GPT-4o ─────────────────────────────────────────────────────────── -GPT_MODEL = "gpt-4o" -GPT_API_URL = "https://api.openai.com/v1/chat/completions" -#GPT_MAX_TOKENS = 16384 - -# ── ChromaDB ───────────────────────────────────────────────────────────────── -CHROMA_PERSIST_DIR = os.getenv("CHROMA_PERSIST_DIR", "./chroma_store") -CHROMA_COLLECTION = "doctodb_chunks" -EMBED_MODEL_NAME = "all-MiniLM-L6-v2" -NEIGHBOR_OVERLAP = 0 -# Pre-warm the embedding model at import time so the first PDF doesn't pay -# the cold-load cost. Runs in a daemon thread — doesn't block startup. -def _prewarm_embed_model() -> None: - try: - _get_embed_model() - log.info("Embedding model pre-warmed.") - except Exception as e: - log.warning(f"Embed model pre-warm failed: {e}") - -# ── Batching ───────────────────────────────────────────────────────────────── - -# ── Adaptive config (probed at startup, falls back to safe defaults) ────── -_GEMINI_DEFAULTS = {"max_input_chars": 60_000, "batch_delay": 5, "max_output_tokens": 32_768} -_GPT_DEFAULTS = {"max_input_chars": 50_000, "batch_delay": 2, "max_tokens": 4096} -_config_lock = threading.Lock() -_GEMINI_CONFIG = dict(_GEMINI_DEFAULTS) -_GPT_CONFIG = dict(_GPT_DEFAULTS) - - -def _probe_gemini_limits(model_name: str) -> None: - """Fetch real token limits from Gemini model metadata and update _GEMINI_CONFIG.""" - try: - url = ( - f"https://generativelanguage.googleapis.com/v1beta/models/{model_name}" - f"?key={GEMINI_API_KEY}" - ) - r = requests.get(url, timeout=10) - if r.status_code != 200: - return - info = r.json() - input_tok = info.get("inputTokenLimit", 200_000) - output_tok = info.get("outputTokenLimit", 32_768) - rpm = info.get("rpm", 15) # not always present; default free-tier - - # Use 70% of input limit, leaving headroom for prompt overhead - max_chars = min(int(input_tok * 3 * 0.70), 500_000) - with _config_lock: - _GEMINI_CONFIG["max_input_chars"] = max_chars - _GEMINI_CONFIG["max_output_tokens"] = min(output_tok, 32768) - _GEMINI_CONFIG["batch_delay"] = max(2, round(60 / rpm)) - log.info("Gemini limits probed: input=%d tokens → %d chars, delay=%ds", - input_tok, max_chars, _GEMINI_CONFIG["batch_delay"]) - except Exception as exc: - log.warning("Gemini limit probe failed (%s) — using defaults", exc) - - -def _probe_gpt_limits() -> None: - """Make a 1-token dummy call to GPT-4o and read rate-limit headers.""" - try: - headers = { - "Authorization": f"Bearer {OPENAI_API_KEY}", - "Content-Type": "application/json", - } - body = { - "model": "gpt-4o", - "max_tokens": 1, - "messages": [{"role": "user", "content": "hi"}], - } - r = requests.post( - "https://api.openai.com/v1/chat/completions", - headers=headers, json=body, timeout=15 - ) - # Headers are present even on successful responses - tpm = int(r.headers.get("x-ratelimit-limit-tokens", 30_000)) - rpm = int(r.headers.get("x-ratelimit-limit-requests", 500)) - - # 50% of TPM for input (rest is output budget), 15% of TPM for max_tokens - max_chars = int(tpm * 0.50 * 3) # tokens → chars - max_tok = int(tpm * 0.15) - - with _config_lock: - _GPT_CONFIG["max_input_chars"] = max_chars - _GPT_CONFIG["max_tokens"] = max_tok - _GPT_CONFIG["batch_delay"] = max(1, round(60 / rpm)) - log.info("GPT limits probed: tpm=%d → %d chars, rpm=%d, delay=%ds", - tpm, max_chars, rpm, _GPT_CONFIG["batch_delay"]) - except Exception as exc: - log.warning("GPT limit probe failed (%s) — using defaults", exc) - - -def _run_startup_probes(model_name: str) -> None: - """Run both probes in parallel — called once after model selection.""" - with ThreadPoolExecutor(max_workers=2) as ex: - ex.submit(_probe_gemini_limits, model_name) - ex.submit(_probe_gpt_limits) - - -# MAX_BATCH_CHARS is the shared floor — set after probes complete -def _get_max_batch_chars() -> int: - return min(_GEMINI_CONFIG["max_input_chars"], _GPT_CONFIG["max_input_chars"]) -GEMINI_RETRY_DELAYS = [60, 120, 180] -GPT_RETRY_DELAYS = [30, 60, 90] - -# ── Extraction ──────────────────────────────────────────────────────────────── -MIN_TABLE_ROWS = 2 -MIN_TABLE_COLS = 2 -MIN_CELL_CHARS = 2 -MIN_TEXT_CHARS = 40 -OCR_THRESHOLD = 50 - -# ── Consensus ──────────────────────────────────────────────────────────────── -CONSENSUS_VALUE_TOL = 0.05 # 5 % numeric tolerance for value matching - -CACHE_FILE = "pdf_extraction_cache.json" - -# ───────────────────────────────────────────────────────────────────────────── -# SINGLE RETRIEVAL QUERY -# ───────────────────────────────────────────────────────────────────────────── +GEMINI_MODEL = os.getenv("GEMINI_MODEL", "gemini-2.5-flash") +GEMINI_API_URL = f"https://generativelanguage.googleapis.com/v1beta/models/{GEMINI_MODEL}:generateContent" -RETRIEVAL_QUERY = ( - "Material property data including section category, property name, " - "measured value, SI unit, imperial unit, test condition standard " - "such as ASTM ISO DIN, and comments. Properties include mechanical " - "thermal electrical physical rheological optical categories. " - "Values in MPa GPa percent density conductivity temperature modulus " - "strength elongation hardness viscosity flammability. " - "Material name, manufacturer, trade grade, abbreviation. " - "DOI digital object identifier paper reference link." -) +EMBED_MODEL_NAME = "allenai/scibert_scivocab_uncased" # matches DocToDB_eval_v2 (no-rank pipeline) +VALUE_REL_TOL = 0.005 # 0.5% relative tolerance for numeric matching +MEANING_THRESHOLD = 0.35 # cosine similarity floor for property/material meaning check # ───────────────────────────────────────────────────────────────────────────── -# GEMINI SCHEMA + PROMPT +# SCHEMA / PROMPT # ───────────────────────────────────────────────────────────────────────────── SCHEMA = { "type": "OBJECT", "properties": { - "material_name": {"type": "STRING"}, + "material_name": {"type": "STRING"}, "material_abbreviation": {"type": "STRING"}, - #"trade_grade": {"type": "STRING", "maxLength": 100}, - "manufacturer": {"type": "STRING", "maxLength": 100}, - "doi": {"type": "STRING"}, "mechanical_properties": { "type": "ARRAY", "items": { "type": "OBJECT", "properties": { - "section": {"type": "STRING"}, - "property_name": {"type": "STRING"}, - "value": {"type": "STRING"}, - "unit": {"type": "STRING"}, - "english": {"type": "STRING"}, + "section": {"type": "STRING"}, + "property_name": {"type": "STRING"}, + "value": {"type": "STRING"}, + "unit": {"type": "STRING"}, + "english": {"type": "STRING"}, "test_condition": {"type": "STRING"}, - "comments": {"type": "STRING"}, - "material_name": {"type": "STRING"}, - "source_page": {"type": "STRING"}, - "chunk_type": {"type": "STRING"}, - "source_text": {"type": "STRING"}, + "comments": {"type": "STRING"}, }, - "required": [ - "section", "property_name", "value", "unit", - "english", "test_condition", "comments", - "material_name", "source_page", "chunk_type", - "source_text", - ], + "required": ["section", "property_name", "value", "english", "comments"], }, }, }, } -# EXTRACTION_PROMPT = ( -# "You are an expert materials scientist. " -# "The content below is extracted from a materials datasheet or research paper. " -# "Each block is tagged with its page number and type (TABLE or TEXT).\n\n" -# "Extract every material property. For each property record:\n" -# " - section : category (Mechanical, Thermal, Electrical, Physical, Rheological, etc.)\n" -# " - property_name : exact name as written\n" -# " - value : exact value or range\n" -# " - unit : SI unit\n" -# " - english : imperial equivalent if shown, else ''\n" -# " - test_condition : ASTM/ISO/DIN standard or conditions, else ''\n" -# " - comments : footnotes or qualifications, else ''\n" -# " - material_name : exact material name (never leave blank)\n" -# " - source_page : page number from the block header (digits only, e.g. '3')\n" -# " - chunk_type : 'table' if from a TABLE block, 'text' if from a TEXT block\n\n" -# "Also extract at the top level:\n" -# " - doi : the DOI of the paper/datasheet (e.g. '10.1016/j.polymer.2023.01.001'). " -# " If a clickable 'click here' link or https://doi.org/... URL appears, " -# " extract just the DOI identifier (without the https://doi.org/ prefix). " -# " Leave empty string '' if not found.\n\n" -# "RULES:\n" -# " - Extract ONLY measured/specified properties, NOT equations or simulation params.\n" -# " - If multiple materials appear, create separate entries for each.\n" -# " - Preserve source_page exactly from the block header.\n" -# " - Respond ONLY with valid JSON matching the schema.\n" -# "\n\nCONTENT:\n" -# ) EXTRACTION_PROMPT = ( - "You are an expert materials scientist and information extraction engine.\n" - "The content below is extracted from a materials datasheet or research paper.\n" - "Each block is tagged with its page number and type (TABLE or TEXT).\n\n" - "Extract EVERY material property with a numeric value. This includes properties " - "found in tables, prose sentences, captions, and comparative statements.\n" - "Example: 'The tensile strength of PLA/WF composites was 45.3 MPa' is a valid extraction.\n\n" - "Extract the following TOP-LEVEL fields (one value for the whole document):\n" - " - material_name : material name only, not a full description paragraph.\n" - " - material_abbreviation : short abbreviation (e.g. 'PLA', 'WF/PLA').\n" - " - trade_grade : ONLY a short grade code under 20 characters (e.g. 'PLA002', 'Grade A'). " - " If no specific product grade code exists in the text, leave exactly ''. " - " DO NOT put descriptions, study context, sentences, or multiple grades here. " - " If multiple grades exist, pick the primary one only.\n" - " - manufacturer : company name only (e.g. 'Kuraray Company'). " - "Do not include descriptions or test conditions here.\n" - " - doi : DOI identifier only (e.g. '10.1016/j.polymer.2023.01.001'). " - "Strip any https://doi.org/ prefix. Leave '' if not found.\n\n" - "For each property record extract:\n" - " - section : must be one of: Mechanical, Thermal, Electrical, " - "Physical, Rheological, Optical, Chemical, Barrier, Morphological, Surface, Other\n" - " - property_name : exact name as written in the source\n" - " - value : exact value or midpoint of range\n" - " - unit : unit exactly as written\n" - " - english : imperial equivalent if shown, else ''\n" - " - test_condition : ASTM/ISO/DIN standard or test conditions, else ''\n" - " - comments : footnotes, qualifications, or ± uncertainty, else ''\n" - " - material_name : exact material name for this property (never leave blank)\n" - " - source_page : page number from the block header (digits only, e.g. '3')\n" - " - chunk_type : 'table' if from a TABLE block, 'text' if from a TEXT block\n\n" - " - source_text : copy the exact sentence or table row this property " - "was extracted from. Never leave blank if the property came from text.\n" - "INCLUDE ALL PROPERTY TYPES — not just mechanical:\n" - "tensile strength, modulus, elongation, impact strength, hardness, density, " - "glass transition temperature, melting temperature, decomposition temperature, " - "thermal conductivity, crystallinity, molecular weight, melt flow index, " - "degradation rate, water absorption, electrical conductivity, viscosity, " - "contact angle, and any other numeric material characteristic.\n\n" - "RULES:\n" - " - Extract from BOTH tables and prose — never skip text blocks.\n" - " - Extract values explicitly stated in the text — do not invent values not present.\n" - " - If multiple materials appear, create separate entries for each.\n" - " - NEVER return an empty properties array if any numeric value exists.\n" - " - Preserve source_page exactly from the block header.\n" - " - Respond ONLY with valid JSON matching the schema.\n" - "\n\nCONTENT:\n" -) -# GPT-4o uses plain-text prompt (no structured output schema enforcement, -# but we instruct it to return the same JSON shape) -# GPT_SYSTEM_PROMPT = ( -# "You are an expert materials scientist. " -# "Extract material properties from the provided PDF content blocks. " -# "Respond ONLY with valid JSON in exactly this shape — no markdown, no extra text:\n" -# "{\n" -# ' "material_name": "...",\n' -# ' "material_abbreviation": "...",\n' -# ' "trade_grade": "...",\n' -# ' "manufacturer": "...",\n' -# ' "doi": "...",\n' -# ' "mechanical_properties": [\n' -# ' {\n' -# ' "section": "Mechanical",\n' -# ' "property_name": "Tensile Strength",\n' -# ' "value": "85",\n' -# ' "unit": "MPa",\n' -# ' "english": "",\n' -# ' "test_condition": "ASTM D638",\n' -# ' "comments": "",\n' -# ' "material_name": "ABS",\n' -# ' "source_page": "3",\n' -# ' "chunk_type": "table"\n' -# ' }\n' -# ' ]\n' -# "}\n\n" -# "Rules:\n" -# "- section: Mechanical / Thermal / Electrical / Physical / Rheological / Optical\n" -# "- Extract ONLY measured/specified properties, NOT equations or simulation params.\n" -# "- doi: extract DOI identifier only (e.g. '10.1016/j.mat.2023.01.001'). " -# " If a 'click here' link or https://doi.org/... URL appears in the text, strip the prefix.\n" -# "- source_page: digits only from block header.\n" -# "- chunk_type: 'table' or 'text' from block header.\n" -# "- Leave fields as '' if unknown." -# ) -GPT_SYSTEM_PROMPT = ( - "You are an expert materials-science information extraction engine.\n" - "Your task is to exhaustively extract ALL material-property data from scientific PDF text.\n\n" - "The PDF may contain narrative prose, tables, captions, figure descriptions, " - "comparative statements, and ranges. Extract properties from ALL of them.\n\n" - "OUTPUT RULES:\n" - "- Return ONLY valid JSON. No markdown, no explanations, no comments, no trailing commas.\n" - "- Output MUST exactly match the schema below.\n" - "- Never omit required fields. Use \"\" for unknown or missing values.\n\n" - "- value: ALWAYS fill this. Put the exact number here (e.g. '51.8'). " - "Never leave value blank if a number exists for this property. " - "For ranges like '40–52': put '40–52' in value AND '40' in min_value AND '52' in max_value.\n" - "JSON SCHEMA:\n" + "You are an expert materials scientist. From the attached PDF, extract the material name, " + "abbreviation, and ALL properties across categories (Mechanical, Thermal, Electrical, Physical, " + "Optical, Rheological, etc.). Return them as 'mechanical_properties' (a single list). " + "For each property, you MUST extract:\n" + "- section (category)\n- property_name\n- value (or range)\n- unit\n" + "- english (converted or alternate units, e.g., psi, °F, inches; write '' if not provided)\n" + "- test_condition\n- comments (include any notes, footnotes, standards, remarks; write '' if none)\n" + "All fields including english and comments are REQUIRED. Respond ONLY with valid JSON following the " + "schema below — no markdown, no code fences, no commentary.\n\n" + "SCHEMA:\n" "{\n" ' "material_name": "",\n' ' "material_abbreviation": "",\n' - ' "trade_grade": "",\n' - ' "manufacturer": "",\n' - ' "doi": "",\n' - ' "properties": [\n' - ' {\n' - ' "section": "",\n' - ' "property_name": "",\n' - ' "value": "",\n' - ' "min_value": "",\n' - ' "max_value": "",\n' - ' "unit": "",\n' - ' "test_condition": "",\n' - ' "comments": "",\n' - ' "source_text": "",\n' - ' "material_name": "",\n' - ' "source_page": "",\n' - ' "chunk_type": ""\n' - ' }\n' - ' ]\n' - "}\n\n" - "CRITICAL EXTRACTION RULES:\n" - "- Extract EVERY material property with a numeric value and unit.\n" - "- Extract from BOTH prose and tables — sentences like 'The tensile strength was 45.3 MPa' " - "are valid extractions.\n" - "- Extract ranges and ± values: put the midpoint or reported value in 'value', " - "lower bound in 'min_value', upper bound in 'max_value'.\n" - "- Extract repeated measurements only once unless conditions differ.\n" - "- NEVER ignore a property because it is outside mechanical testing.\n" - "- NEVER return an empty properties array if any numeric material property exists.\n\n" - "INCLUDE ALL PROPERTY TYPES:\n" - "Mechanical, Thermal, Electrical, Physical, Rheological, Optical, " - "Chemical, Barrier, Morphological, Surface, Crystallinity, Biodegradation, Other\n\n" - "EXAMPLES OF VALID PROPERTIES:\n" - "tensile strength, elastic modulus, Young's modulus, flexural modulus, " - "elongation at break, impact strength, hardness, density, melt flow index, " - "viscosity, glass transition temperature, melting temperature, decomposition temperature, " - "thermal conductivity, electrical conductivity, crystallinity, molecular weight, " - "degradation rate, water absorption, contact angle\n\n" - "FIELD RULES:\n" - "- section must be one of: Mechanical, Thermal, Electrical, Physical, Rheological, " - "Optical, Chemical, Barrier, Morphological, Surface, Other\n" - "- source_page: digits only from the block header.\n" - "- chunk_type: 'table' or 'text' from the block header.\n" - "- source_text: copy the exact sentence or table row this property was extracted from.\n" - "- material_name: exact material name, never leave blank.\n" - "- Leave all other fields as \"\" if unknown — never omit them.\n\n" - "NORMALIZATION RULES:\n" - "- Preserve original units and property names exactly as written.\n" - "- Do NOT convert units, infer missing values, or hallucinate standards.\n" - "- If multiple materials appear, extract the primary material being characterized.\n\n" - "ANTI-HALLUCINATION:\n" - "- Only extract values explicitly stated in the provided text.\n" - "- If a value is not present, do not invent it.\n" - "- Missing a weakly-labeled property is acceptable. Inventing one is not.\n" -) -# ───────────────────────────────────────────────────────────────────────────── -# BOILERPLATE DETECTION -# ───────────────────────────────────────────────────────────────────────────── - -_SKIP_HEADING_RE = re.compile( - r"^(references|bibliography|acknowledgements?|table\s+of\s+contents|" - r"copyright|legal\s+notice|disclaimer|index|appendix\s+[a-z]$)", - re.IGNORECASE, + ' "mechanical_properties": [\n' + ' {"section":"","property_name":"","value":"","unit":"","english":"","test_condition":"","comments":""}\n' + " ]\n" + "}\n" ) -def _is_boilerplate(text: str) -> bool: - return bool(_SKIP_HEADING_RE.match(text.strip().split("\n")[0].strip())) - -# ───────────────────────────────────────────────────────────────────────────── -# DATA STRUCTURES -# ───────────────────────────────────────────────────────────────────────────── - -@dataclass -class Chunk: - page_num: int - chunk_type: str - source: str - raw_rows: Optional[List[List[str]]] = None - raw_text: Optional[str] = None - text: str = field(init=False) - score: float = 0.0 - relevant: bool = False - - def __post_init__(self): - if self.chunk_type == "table" and self.raw_rows: - self.text = _rows_to_text(self.raw_rows) - elif self.raw_text: - self.text = self.raw_text.strip() - else: - self.text = "" - - @property - def block_header(self) -> str: - return ( - f"\n\n{'─'*60}\n" - f"[{self.chunk_type.upper()} | Page {self.page_num} | " - f"score={self.score:.3f}]\n" - f"{'─'*60}\n" - ) - -# ───────────────────────────────────────────────────────────────────────────── -# HELPERS -# ───────────────────────────────────────────────────────────────────────────── - -def _rows_to_text(rows: List[List[Any]]) -> str: - lines = [] - for row in rows: - cells = [str(c).strip() if c is not None else "" for c in row] - if any(len(c) >= MIN_CELL_CHARS for c in cells): - lines.append(" | ".join(cells)) - return _fix_spaced_text("\n".join(lines)) - - -def _is_valid_table(rows: List[List[Any]]) -> bool: - if not rows or len(rows) < MIN_TABLE_ROWS: - return False - if max((len(r) for r in rows), default=0) < MIN_TABLE_COLS: - return False - for row in rows[1:]: - if any(re.search(r"\d", str(c)) for c in row): - return True - return False - - -def _split_paragraphs(raw_text: str) -> List[str]: - paragraphs: List[str] = [] - current: List[str] = [] - for line in raw_text.splitlines(): - stripped = line.strip() - if not stripped: - if current: - paragraphs.append(" ".join(current)) - current = [] - else: - current.append(stripped) - if current: - paragraphs.append(" ".join(current)) - return [_fix_spaced_text(p) for p in paragraphs if len(p) >= MIN_TEXT_CHARS] - - -def _normalise_doi(raw: str) -> str: - """Strip URL prefix and whitespace, return bare DOI or empty string.""" - if not raw: - return "" - doi = raw.strip() - for prefix in ("https://doi.org/", "http://doi.org/", "doi.org/", "DOI:", "doi:"): - if doi.lower().startswith(prefix.lower()): - doi = doi[len(prefix):] - return doi.strip() - - -def _doi_url(doi: str) -> str: - """Return full clickable URL for a bare DOI.""" - doi = _normalise_doi(doi) - if not doi: - return "" - return f"https://doi.org/{doi}" - -# ───────────────────────────────────────────────────────────────────────────── -# CACHE -# ───────────────────────────────────────────────────────────────────────────── - -def _pdf_hash(pdf_bytes: bytes) -> str: - return hashlib.sha256(pdf_bytes).hexdigest()[:16] - -def _load_cache() -> Dict: - try: - if os.path.exists(CACHE_FILE): - with open(CACHE_FILE) as f: - return json.load(f) - except Exception: - pass - return {} - -def _save_cache(cache: Dict): - try: - with open(CACHE_FILE, "w") as f: - json.dump(cache, f) - except Exception as e: - log.warning(f"Cache save failed: {e}") - -def cache_get(pdf_bytes: bytes) -> Optional[Dict]: - return _load_cache().get(_pdf_hash(pdf_bytes)) - -def cache_set(pdf_bytes: bytes, result: Dict): - cache = _load_cache() - cache[_pdf_hash(pdf_bytes)] = result - _save_cache(cache) - -# ───────────────────────────────────────────────────────────────────────────── -# CHROMADB CLIENT -# ───────────────────────────────────────────────────────────────────────────── - -_chroma_client: Optional[Any] = None -_chroma_collection: Optional[Any] = None -_embed_model: Optional[Any] = None - - -def _get_embed_model() -> Any: - global _embed_model - if _embed_model is None: - if not ST_AVAILABLE: - raise ImportError("sentence-transformers not installed.\nRun: pip install sentence-transformers") - log.info(f"Loading embedding model '{EMBED_MODEL_NAME}' …") - _embed_model = SentenceTransformer(EMBED_MODEL_NAME) - log.info("Embedding model ready.") - return _embed_model - -threading.Thread(target=_prewarm_embed_model, daemon=True).start() - -def _get_chroma_collection() -> Any: - global _chroma_client, _chroma_collection - if _chroma_collection is not None: - return _chroma_collection - if not CHROMA_AVAILABLE: - raise ImportError("chromadb not installed. Run: pip install chromadb") - os.makedirs(CHROMA_PERSIST_DIR, exist_ok=True) - _chroma_client = chromadb.PersistentClient(path=CHROMA_PERSIST_DIR) - _chroma_collection = _chroma_client.get_or_create_collection( - name=CHROMA_COLLECTION, - metadata={"hnsw:space": "cosine"}, - ) - log.info(f"ChromaDB ready — '{CHROMA_COLLECTION}' ({_chroma_collection.count()} existing vectors)") - return _chroma_collection - - -def _chroma_pdf_exists(pdf_hash: str) -> bool: - try: - col = _get_chroma_collection() - res = col.get(where={"pdf_hash": pdf_hash}, limit=1) - return len(res["ids"]) > 0 - except Exception: - return False - - -def _chroma_store_chunks(chunks: List[Chunk], pdf_hash: str) -> None: - col = _get_chroma_collection() - model = _get_embed_model() - texts = [c.text for c in chunks] - metadatas = [ - {"pdf_hash": pdf_hash, "page_num": c.page_num, - "chunk_type": c.chunk_type, "source": c.source} - for c in chunks - ] - ids = [f"{pdf_hash}_{i}" for i in range(len(chunks))] - # all_embeddings: List[List[float]] = [] - # for start in range(0, len(texts), 64): - # batch = texts[start : start + 64] - # vecs = model.encode(batch, normalize_embeddings=True) - # all_embeddings.extend(vecs.tolist()) - vecs = model.encode( - texts, - normalize_embeddings=True, - batch_size=256, # internal mini-batch size for GPU memory - show_progress_bar=False, - ) - all_embeddings = vecs.tolist() - - col.upsert(ids=ids, documents=texts, embeddings=all_embeddings, metadatas=metadatas) - log.info(f"ChromaDB: stored {len(chunks)} chunks for pdf_hash={pdf_hash}") - - -def _chroma_rank_all(pdf_hash: str) -> List[Chunk]: - col = _get_chroma_collection() - model = _get_embed_model() - all_ids = col.get(where={"pdf_hash": pdf_hash}, include=[]) - n_total = len(all_ids["ids"]) - if n_total == 0: - log.warning(f"No chunks in ChromaDB for pdf_hash={pdf_hash}") - return [] - query_vec = model.encode([RETRIEVAL_QUERY], normalize_embeddings=True)[0].tolist() - results = col.query( - query_embeddings=[query_vec], - n_results=n_total, - where={"pdf_hash": pdf_hash}, - include=["documents", "metadatas", "distances"], - ) - ranked_chunks: List[Chunk] = [] - for doc, meta, dist in zip( - results.get("documents", [[]])[0], - results.get("metadatas", [[]])[0], - results.get("distances", [[]])[0], - ): - similarity = 1.0 - float(dist) - chunk = Chunk( - page_num=int(meta.get("page_num", 1)), - chunk_type=meta.get("chunk_type", "text"), - source=meta.get("source", ""), - raw_text=doc, - score=similarity, - relevant=True, - ) - ranked_chunks.append(chunk) - table_n = sum(1 for c in ranked_chunks if c.chunk_type == "table") - text_n = sum(1 for c in ranked_chunks if c.chunk_type == "text") - log.info(f"ChromaDB ranked {len(ranked_chunks)} chunks ({table_n} tables + {text_n} text)") - return ranked_chunks - -# ───────────────────────────────────────────────────────────────────────────── -# EXTRACTION — TABLES -# ───────────────────────────────────────────────────────────────────────────── - -def _extract_tables_docling(pdf_bytes: bytes) -> List[Chunk]: - chunks: List[Chunk] = [] - if not DOCLING_AVAILABLE: - return chunks - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: - tmp.write(pdf_bytes); tmp_path = tmp.name - try: - converter = DocumentConverter() - result = converter.convert(tmp_path) - doc = result.document - for table in doc.tables: - rows = [[cell.text for cell in row] for row in table.data.grid] - page_num = table.prov[0].page_no if table.prov else 1 - if _is_valid_table(rows): - chunks.append(Chunk(page_num=page_num, chunk_type="table", source="docling", raw_rows=rows)) - log.info(f"Docling tables: {len(chunks)}") - except Exception as e: - log.error(f"Docling table extraction failed: {e}") - finally: - os.unlink(tmp_path) - return chunks - - -def _extract_tables_pdfplumber(pdf_bytes: bytes) -> List[Chunk]: - chunks: List[Chunk] = [] - try: - with pdfplumber.open(io.BytesIO(pdf_bytes)) as pdf: - for page_num, page in enumerate(pdf.pages, start=1): - for strategy in ( - {"vertical_strategy": "lines_strict", "horizontal_strategy": "lines_strict", - "snap_tolerance": 3, "join_tolerance": 3}, - {"vertical_strategy": "text", "horizontal_strategy": "text"}, - ): - tables = page.extract_tables(table_settings=strategy) or [] - for rows in tables: - cleaned = [[str(c).strip() if c else "" for c in row] for row in rows] - if _is_valid_table(cleaned): - chunks.append(Chunk(page_num=page_num, chunk_type="table", - source="pdfplumber", raw_rows=cleaned)) - if tables: - break - log.info(f"pdfplumber tables: {len(chunks)}") - except Exception as e: - log.error(f"pdfplumber table extraction failed: {e}") - return chunks - - -def _extract_tables_camelot(pdf_bytes: bytes) -> List[Chunk]: - chunks: List[Chunk] = [] - if not CAMELOT_AVAILABLE: - return chunks - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: - tmp.write(pdf_bytes); tmp_path = tmp.name - try: - for flavor in ("lattice", "stream"): - try: - tables = camelot.read_pdf(tmp_path, pages="all", flavor=flavor) - if flavor == "stream": - tables = [t for t in tables if t.parsing_report.get("accuracy", 0) > 65] - for table in tables: - rows = [[str(c).strip() for c in row] for row in table.df.values.tolist()] - if _is_valid_table(rows): - chunks.append(Chunk(page_num=table.page, chunk_type="table", - source=f"camelot-{flavor}", raw_rows=rows)) - if chunks: - break - except Exception as e: - log.warning(f"camelot {flavor}: {e}") - finally: - os.unlink(tmp_path) - log.info(f"camelot tables: {len(chunks)}") - return chunks - -# ───────────────────────────────────────────────────────────────────────────── -# EXTRACTION — TEXT -# ───────────────────────────────────────────────────────────────────────────── - -def _extract_text_docling(pdf_bytes: bytes) -> List[Chunk]: - chunks: List[Chunk] = [] - if not DOCLING_AVAILABLE: - return chunks - with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: - tmp.write(pdf_bytes); tmp_path = tmp.name - try: - converter = DocumentConverter() - result = converter.convert(tmp_path) - doc = result.document - for item, _ in doc.iterate_items(): - text = getattr(item, "text", "").strip() - page_num = item.prov[0].page_no if item.prov else 1 - if not text or len(text) < MIN_TEXT_CHARS or _is_boilerplate(text): - continue - chunks.append(Chunk(page_num=page_num, chunk_type="text", source="docling", raw_text=text)) - log.info(f"Docling text chunks: {len(chunks)}") - except Exception as e: - log.error(f"Docling text extraction failed: {e}") - finally: - os.unlink(tmp_path) - return chunks - - -def _extract_text_pymupdf(pdf_bytes: bytes) -> List[Chunk]: - chunks: List[Chunk] = [] - try: - with fitz.open(stream=pdf_bytes, filetype="pdf") as doc: - for page_idx, page in enumerate(doc): - page_num = page_idx + 1 - raw = page.get_text("text", sort=True) or "" - for para in _split_paragraphs(raw): - if not _is_boilerplate(para): - chunks.append(Chunk(page_num=page_num, chunk_type="text", - source="pymupdf", raw_text=para)) - log.info(f"pymupdf text chunks: {len(chunks)}") - except Exception as e: - log.error(f"pymupdf text extraction failed: {e}") - return chunks - - -def _extract_text_pdfplumber(pdf_bytes: bytes) -> List[Chunk]: - chunks: List[Chunk] = [] - try: - with pdfplumber.open(io.BytesIO(pdf_bytes)) as pdf: - for page_num, page in enumerate(pdf.pages, start=1): - raw = page.extract_text(x_tolerance=3, y_tolerance=3) or "" - for para in _split_paragraphs(raw): - if not _is_boilerplate(para): - chunks.append(Chunk(page_num=page_num, chunk_type="text", - source="pdfplumber-text", raw_text=para)) - log.info(f"pdfplumber text chunks: {len(chunks)}") - except Exception as e: - log.error(f"pdfplumber text extraction failed: {e}") - return chunks - - -def _ocr_page(pdf_bytes: bytes, page_num: int) -> str: - if not OCR_AVAILABLE: - return "" - try: - with fitz.open(stream=pdf_bytes, filetype="pdf") as doc: - page = doc[page_num - 1] - mat = fitz.Matrix(300 / 72, 300 / 72) - pix = page.get_pixmap(matrix=mat) - img = Image.open(io.BytesIO(pix.tobytes("png"))) - return pytesseract.image_to_string(img, lang="eng") or "" - except Exception as e: - log.warning(f"OCR page {page_num}: {e}") - return "" - - -def _verify_page_coverage(pdf_bytes: bytes, chunks: List[Chunk]) -> List[Chunk]: - with fitz.open(stream=pdf_bytes, filetype="pdf") as doc: - total_pages = set(range(1, len(doc) + 1)) - covered = {c.page_num for c in chunks} - missing = total_pages - covered - if not missing: - return chunks - log.warning(f"Page coverage: {len(missing)} pages missing — OCRing: {sorted(missing)}") - for page_num in sorted(missing): - ocr_text = _ocr_page(pdf_bytes, page_num) - for para in _split_paragraphs(ocr_text): - if not _is_boilerplate(para): - chunks.append(Chunk(page_num=page_num, chunk_type="text", - source="ocr-fallback", raw_text=para)) - return chunks - -def _dedup(chunks: List[Chunk]) -> List[Chunk]: - seen: set = set() - unique: List[Chunk] = [] - for c in chunks: - key = (c.chunk_type, re.sub(r"\s+", " ", c.text.strip())[:200]) - if key not in seen: - seen.add(key) - unique.append(c) - return unique +def make_abbreviation(name: str) -> str: + if not name: + return "UNKNOWN" + words = name.split() + abbr = "".join(w[0] for w in words if w and w[0].isalpha()).upper() + return abbr or name[:6].upper() -# def extract_all_chunks(pdf_bytes: bytes) -> List[Chunk]: -# table_chunks: List[Chunk] = [] -# text_chunks: List[Chunk] = [] -# docling_tables = _extract_tables_docling(pdf_bytes) -# table_chunks.extend(docling_tables) -# table_chunks.extend(_extract_tables_pdfplumber(pdf_bytes)) -# table_chunks.extend(_extract_tables_camelot(pdf_bytes)) -# table_chunks = _dedup(table_chunks) -# if DOCLING_AVAILABLE and docling_tables: -# text_chunks.extend(_extract_text_docling(pdf_bytes)) -# pymupdf_text = _extract_text_pymupdf(pdf_bytes) -# text_chunks.extend(pymupdf_text) -# if sum(len(c.text) for c in pymupdf_text) < 500: -# log.warning("pymupdf sparse — supplementing with pdfplumber") -# text_chunks.extend(_extract_text_pdfplumber(pdf_bytes)) -# text_chunks = _dedup(text_chunks) -# all_chunks = table_chunks + text_chunks -# all_chunks = _verify_page_coverage(pdf_bytes, all_chunks) -# log.info(f"Total chunks: {len(all_chunks)} ({len(table_chunks)} tables + {len(text_chunks)} text)") -# return all_chunks - -def extract_all_chunks(pdf_bytes: bytes) -> List[Chunk]: - with ThreadPoolExecutor(max_workers=5) as ex: - futures = { - ex.submit(_extract_tables_docling, pdf_bytes): "docling_tables", - ex.submit(_extract_tables_pdfplumber, pdf_bytes): "pdfplumber_tables", - ex.submit(_extract_tables_camelot, pdf_bytes): "camelot_tables", - ex.submit(_extract_text_pymupdf, pdf_bytes): "pymupdf_text", - ex.submit(_extract_text_pdfplumber, pdf_bytes): "pdfplumber_text", - } - results = {} - for future in as_completed(futures): - name = futures[future] - try: - results[name] = future.result() - except Exception as e: - log.error(f"{name} failed: {e}") - results[name] = [] - - docling_tables = results.get("docling_tables", []) - pdfplumber_tables= results.get("pdfplumber_tables", []) - camelot_tables = results.get("camelot_tables", []) - pymupdf_text = results.get("pymupdf_text", []) - pdfplumber_text = results.get("pdfplumber_text", []) - - table_chunks = _dedup(docling_tables + pdfplumber_tables + camelot_tables) - - text_chunks: List[Chunk] = [] - if DOCLING_AVAILABLE and docling_tables: - # docling text runs after we know docling tables succeeded - text_chunks.extend(_extract_text_docling(pdf_bytes)) - text_chunks.extend(pymupdf_text) - if sum(len(c.text) for c in pymupdf_text) < 500: - log.warning("pymupdf sparse — supplementing with pdfplumber text") - text_chunks.extend(pdfplumber_text) - text_chunks = _dedup(text_chunks) - - all_chunks = table_chunks + text_chunks - all_chunks = _verify_page_coverage(pdf_bytes, all_chunks) - log.info( - f"Total chunks: {len(all_chunks)} " - f"({len(table_chunks)} tables + {len(text_chunks)} text)" - ) - return all_chunks - -def _fix_spaced_text(text: str) -> str: - """Fix PDF kerning artifacts like 'C O M P O S I T E' → 'COMPOSITE'.""" - # Pattern: single chars separated by spaces forming a word - return re.sub( - r'\b([A-Z]) (?=[A-Z] |[A-Z]\b)', - r'\1', - text - ) -# ───────────────────────────────────────────────────────────────────────────── -# CHROMADB INDEX -# ───────────────────────────────────────────────────────────────────────────── - -def _build_overlapping_chunks(chunks: List[Chunk], neighbor_window: int = NEIGHBOR_OVERLAP) -> List[Chunk]: - if neighbor_window <= 0: - return chunks - eligible = [(i, c) for i, c in enumerate(chunks) if not _is_boilerplate(c.text)] - overlap_chunks: List[Chunk] = [] - seen_content: set = set() - for pos, (_, centre) in enumerate(eligible): - if centre.chunk_type == "table": - continue - parts: List[str] = [] - for offset in range(-neighbor_window, neighbor_window + 1): - nb_pos = pos + offset - if nb_pos < 0 or nb_pos >= len(eligible): - continue - _, nb_chunk = eligible[nb_pos] - if nb_chunk.chunk_type == "table": - continue - parts.append(nb_chunk.text.strip()) - merged = " \n\n ".join(p for p in parts if p) - if not merged or merged in seen_content: - continue - seen_content.add(merged) - overlap_chunks.append(Chunk(page_num=centre.page_num, chunk_type="text", - source=f"{centre.source}+overlap", raw_text=merged)) - return chunks + overlap_chunks - - -def index_chunks_in_chroma(chunks: List[Chunk], pdf_hash: str) -> None: - if _chroma_pdf_exists(pdf_hash): - log.info(f"ChromaDB: pdf_hash={pdf_hash} already indexed — skipping.") - return - all_chunks = _build_overlapping_chunks(chunks, neighbor_window=NEIGHBOR_OVERLAP) - storable = [c for c in all_chunks if not _is_boilerplate(c.text) and c.text.strip()] - _chroma_store_chunks(storable, pdf_hash) - -# ───────────────────────────────────────────────────────────────────────────── -# BATCH BUILDER -# ───────────────────────────────────────────────────────────────────────────── - -def build_batches(chunks: List[Chunk]) -> List[str]: - if not chunks: - return [] - batches: List[str] = [] - current_batch: List[Chunk] = [] - current_chars: int = 0 - for chunk in chunks: - entry_len = len(chunk.block_header) + len(chunk.text) - _max = _get_max_batch_chars() - if entry_len > _max: - chunk.text = chunk.text[: _max - len(chunk.block_header) - 20] - entry_len = len(chunk.block_header) + len(chunk.text) - - if current_chars + entry_len > _max and current_batch: - current_batch.sort(key=lambda c: c.page_num) - batches.append("\n".join(c.block_header + c.text for c in current_batch)) - current_batch = [] - current_chars = 0 - current_batch.append(chunk) - current_chars += entry_len - if current_batch: - current_batch.sort(key=lambda c: c.page_num) - batches.append("\n".join(c.block_header + c.text for c in current_batch)) - log.info(f"Built {len(batches)} batch(es) from {len(chunks)} chunks") - return batches - # ───────────────────────────────────────────────────────────────────────────── -# GEMINI EXTRACTION +# GEMINI CALL — single call, whole PDF, no chunking # ───────────────────────────────────────────────────────────────────────────── -def _call_gemini(text_payload: str) -> Tuple[Optional[Dict], str]: - log.info(f"Gemini calling with maxOutputTokens={_GEMINI_CONFIG['max_output_tokens']}") - log.info(f"Gemini input (first 300 chars): {repr(text_payload[:300])}") # add this +def call_gemini_from_bytes(pdf_bytes: bytes) -> Optional[Dict[str, Any]]: + if not GEMINI_API_KEY: + st.error("GEMINI_API_KEY not set.") + return None + encoded = base64.b64encode(pdf_bytes).decode("utf-8") payload = { - "contents": [{"parts": [{"text": EXTRACTION_PROMPT + text_payload}]}], + "contents": [{ + "parts": [ + {"text": EXTRACTION_PROMPT}, + {"inlineData": {"mimeType": "application/pdf", "data": encoded}}, + ] + }], "generationConfig": { "temperature": 0, "responseMimeType": "application/json", "responseSchema": SCHEMA, - "maxOutputTokens": _GEMINI_CONFIG["max_output_tokens"], }, } - for attempt, wait in enumerate([0] + GEMINI_RETRY_DELAYS): - if wait: - log.warning(f"Gemini 429 — waiting {wait}s (retry {attempt})") - time.sleep(wait) - try: - resp = requests.post(GEMINI_API_URL, json=payload, timeout=600) - if resp.status_code == 429: - continue - if not resp.ok: - return None, f"Gemini HTTP {resp.status_code}: {resp.text[:2000]}" - data = resp.json() - candidates = data.get("candidates", []) - if not candidates: - return None, f"Gemini no candidates: {json.dumps(data)[:400]}" - parts = candidates[0].get("content", {}).get("parts", []) - for part in parts: - text = part.get("text", "").strip() - if text.startswith("{"): - try: - # Layer 2 — regex truncation of abused top-level string fields - text = re.sub( - r'("(?:trade_grade|manufacturer)"\s*:\s*")' - r'([^"]{150})[^"]*"', - r'\1\2"', - text - ) - - parsed = json.loads(text) - - # Layer 3 — post-parse semantic validation - tg = parsed.get("trade_grade", "") - if tg and ( - len(tg) > 30 - or tg.count(" ") > 3 - or tg.count("_") > 2 - or any(w in tg.lower() for w in - ["study", "used", "modified", "provided", "analysis", - "document", "without", "unless", "specified", "otherwise", - "percent", "weight", "fiber", "composite", "matrix"]) - ): - log.warning(f"Clearing abusive trade_grade: {repr(tg[:80])}") - parsed["trade_grade"] = "" - - mfr = parsed.get("manufacturer", "") - if mfr and (len(mfr) > 100 or mfr.count("_") > 2): - log.warning(f"Truncating abusive manufacturer: {repr(mfr[:80])}") - parsed["manufacturer"] = mfr[:100].rsplit(" ", 1)[0] - - return parsed, "" - - except json.JSONDecodeError as e: - log.error(f"Gemini raw response (first 500 chars): {repr(text[:500])}") - return None, f"Gemini JSON parse error: {e}" - - return None, "No JSON in Gemini response." - except requests.Timeout: - return None, "Gemini timeout after 600s." - except Exception as e: - return None, f"Gemini unexpected: {e}" - return None, "Gemini max retries exceeded." - - -def run_gemini_batches(batches: List[str], progress_callback=None) -> Tuple[List[Dict], List[str]]: - results: List[Dict] = [] - errors: List[str] = [] - total = len(batches) - for idx, batch_text in enumerate(batches): - pct = 0.60 + 0.17 * (idx / max(total, 1)) - msg = f"[Gemini] batch {idx+1}/{total}…" - log.info(msg) - if progress_callback: - progress_callback(msg, pct) - if idx > 0: - time.sleep(_GEMINI_CONFIG["batch_delay"]) - - result, err = _call_gemini(batch_text) - if result: - results.append(result) - if err: - errors.append(f"Gemini batch {idx+1}: {err}") - log.error(f"Gemini batch {idx+1} error: {err}") - return results, errors - -# ───────────────────────────────────────────────────────────────────────────── -# GPT-4o EXTRACTION -# ───────────────────────────────────────────────────────────────────────────── - -def _call_gpt(text_payload: str) -> Tuple[Optional[Dict], str]: - if not OPENAI_API_KEY: - return None, "OPENAI_API_KEY not set." - headers = { - "Authorization": f"Bearer {OPENAI_API_KEY}", - "Content-Type": "application/json", - } - payload = { - "model": GPT_MODEL, - "temperature": 0, - "response_format": {"type": "json_object"}, - #"max_tokens": GPT_MAX_TOKENS, - "messages": [ - {"role": "system", "content": GPT_SYSTEM_PROMPT}, - {"role": "user", "content": EXTRACTION_PROMPT + text_payload}, - ], - } - for attempt, wait in enumerate([0] + GPT_RETRY_DELAYS): - if wait: - log.warning(f"GPT-4o 429 — waiting {wait}s (retry {attempt})") - time.sleep(wait) - try: - resp = requests.post(GPT_API_URL, headers=headers, json=payload, timeout=300) - if resp.status_code == 429: - continue - if not resp.ok: - return None, f"GPT HTTP {resp.status_code}: {resp.text[:2000]}" - data = resp.json() - content = data.get("choices", [{}])[0].get("message", {}).get("content", "").strip() - # Strip markdown code fences if present - content = re.sub(r"^```(?:json)?\s*", "", content) - content = re.sub(r"\s*```$", "", content) - if content.startswith("{"): - try: - return json.loads(content), "" - except json.JSONDecodeError as e: - return None, f"GPT JSON parse error: {e} | raw: {content[:300]}" - return None, f"GPT response not JSON: {content[:300]}" - except requests.Timeout: - return None, "GPT timeout after 300s." - except Exception as e: - return None, f"GPT unexpected: {e}" - return None, "GPT max retries exceeded." - - -def run_gpt_batches(batches: List[str], progress_callback=None) -> Tuple[List[Dict], List[str]]: - results: List[Dict] = [] - errors: List[str] = [] - total = len(batches) - for idx, batch_text in enumerate(batches): - pct = 0.60 + 0.17 * (idx / max(total, 1)) - msg = f"[GPT-4o] batch {idx+1}/{total}…" - log.info(msg) - if progress_callback: - progress_callback(msg, pct) - if idx > 0: - time.sleep(_GPT_CONFIG["batch_delay"]) - result, err = _call_gpt(batch_text) - if result: - results.append(result) - if err: - errors.append(f"GPT batch {idx+1}: {err}") - log.error(f"GPT batch {idx+1} error: {err}") - return results, errors - -# ───────────────────────────────────────────────────────────────────────────── -# PARALLEL DUAL-LLM RUNNER -# Both LLMs get the same batches, run in parallel threads. -# Each respects its own rate-limit delays — they never share state. -# ───────────────────────────────────────────────────────────────────────────── - -def run_dual_llm_parallel( - batches: List[str], - progress_callback=None, -) -> Tuple[List[Dict], List[Dict], List[str]]: - """ - Returns (gemini_results, gpt_results, all_errors). - Both LLMs run concurrently in separate threads. - """ - gemini_results: List[Dict] = [] - gpt_results: List[Dict] = [] - all_errors: List[str] = [] - - def _gemini_task(): - return run_gemini_batches(batches, progress_callback=None) - - def _gpt_task(): - return run_gpt_batches(batches, progress_callback=None) - - if progress_callback: - progress_callback("Running Gemini + GPT-4o in parallel…", 0.62) - - with ThreadPoolExecutor(max_workers=2) as executor: - future_gemini = executor.submit(_gemini_task) - future_gpt = executor.submit(_gpt_task) - - for future in as_completed([future_gemini, future_gpt]): - if future is future_gemini: - try: - res, errs = future.result() - gemini_results.extend(res) - all_errors.extend(errs) - log.info(f"Gemini done — {len(res)} batch result(s)") - except Exception as e: - all_errors.append(f"Gemini thread error: {e}") - else: - try: - res, errs = future.result() - gpt_results.extend(res) - all_errors.extend(errs) - log.info(f"GPT-4o done — {len(res)} batch result(s)") - except Exception as e: - all_errors.append(f"GPT thread error: {e}") - - return gemini_results, gpt_results, all_errors - -# ───────────────────────────────────────────────────────────────────────────── -# MERGE RESULTS INTO DATAFRAME (per-LLM) -# ───────────────────────────���───────────────────────────────────────────────── - -FIXED_COLS = ["material_name", "material_abbreviation", "manufacturer", "doi"] - - -def _make_abbreviation(name: str) -> str: - abbr = "".join(c for c in name if c.isupper()) - return abbr or name[:4].upper() - - -def _norm(s: str) -> str: - return re.sub(r"[^a-z0-9]", "", str(s).lower().strip()) - - -def _richness(row: Dict) -> int: - return sum(1 for v in row.values() if str(v).strip() not in ("", "N/A")) - - -def _fuzzy_dedup(rows: List[Dict]) -> List[Dict]: - seen_fuzzy: Dict[tuple, int] = {} - kept: List[Dict] = [] - for row in rows: - key = ( - _norm(row.get("material_name", "")), - _norm(row.get("property_name", "")), - _norm(row.get("section", "")), - ) - if key in seen_fuzzy: - existing_idx = seen_fuzzy[key] - existing = kept[existing_idx] - if _richness(row) > _richness(existing): - merged = dict(row) - for k, v in existing.items(): - if str(merged.get(k, "")).strip() in ("", "N/A") and str(v).strip() not in ("", "N/A"): - merged[k] = v - kept[existing_idx] = merged - else: - seen_fuzzy[key] = len(kept) - kept.append(row) - return kept - - -def _extract_doi_from_results(results: List[Dict]) -> str: - """Find the first non-empty DOI across all batch results.""" - for r in results: - doi = _normalise_doi(r.get("doi", "")) - if doi: - return doi - return "" - - -def merge_to_dataframe(results: List[Dict]) -> Tuple[pd.DataFrame, str]: - """ - Returns (DataFrame, doi_string). - """ - all_rows: List[Dict] = [] - seen_exact: set = set() - - fallback = {k: "" for k in FIXED_COLS} - for r in results: - for k in FIXED_COLS: - if not fallback[k] and r.get(k): - fallback[k] = r[k] - if not fallback["material_abbreviation"] and fallback["material_name"]: - fallback["material_abbreviation"] = _make_abbreviation(fallback["material_name"]) - - doi = _normalise_doi(fallback.get("doi", "")) - - for r in results: - r_doi = _normalise_doi(r.get("doi", "")) - if r_doi and not doi: - doi = r_doi - - r_identity = { - "material_name": r.get("material_name", "") or fallback["material_name"], - "material_abbreviation": r.get("material_abbreviation", "") or fallback["material_abbreviation"], - "manufacturer": r.get("manufacturer", "") or fallback["manufacturer"], - } - if not r_identity["material_abbreviation"] and r_identity["material_name"]: - r_identity["material_abbreviation"] = _make_abbreviation(r_identity["material_name"]) - - for item in r.get("properties", r.get("mechanical_properties", [])): - prop_mat = item.get("material_name", "").strip() - identity = dict(r_identity) - if prop_mat: - identity["material_name"] = prop_mat - identity["material_abbreviation"] = _make_abbreviation(prop_mat) - - raw_page = str(item.get("source_page", "")).strip() - page_label = ( - f"Page {raw_page}" if raw_page.isdigit() - else raw_page if raw_page - else "Unknown" - ) - chunk_type = item.get("chunk_type", "").strip() or "unknown" - - key = ( - _norm(identity["material_name"]), - _norm(item.get("section", "")), - _norm(item.get("property_name", "")), - _norm(item.get("value", "")), - ) - if key in seen_exact: - continue - seen_exact.add(key) - - # all_rows.append({ - # **identity, - # "section": item.get("section", "") or "General", - # "property_name": item.get("property_name", "") or "Unknown", - # "value": item.get("value", "") or "N/A", - # "unit": item.get("unit", "") or "", - # "english": item.get("english", "") or "", - # "test_condition": item.get("test_condition", "") or "", - # "comments": item.get("comments", "") or "", - # "source_page": page_label, - # "chunk_type": chunk_type, - # }) - value = item.get("value", "").strip() - min_v = item.get("min_value", "").strip() - max_v = item.get("max_value", "").strip() + try: + r = requests.post(GEMINI_API_URL, params={"key": GEMINI_API_KEY}, json=payload, timeout=300) + if not r.ok: + st.error(f"Gemini HTTP {r.status_code}") + st.code(r.text) + return None + data = r.json() + candidates = data.get("candidates", []) + if not candidates: + st.warning("Gemini returned no candidates.") + return None + for part in candidates[0].get("content", {}).get("parts", []): + text = part.get("text", "").strip() + if text.startswith("{"): + return json.loads(text) + st.warning("Gemini response didn't contain JSON text.") + return None + except Exception as e: + st.error(f"Gemini API exception: {e}") + return None - # if GPT left value blank but filled min/max, reconstruct it - if not value and min_v and max_v: - value = f"{min_v}–{max_v}" - elif not value and min_v: - value = min_v - elif not value and max_v: - value = max_v - all_rows.append({ - **identity, - "section": item.get("section", "") or "General", - "property_name": item.get("property_name", "") or "Unknown", - "value": value or "N/A", - "min_value": min_v, - "max_value": max_v, - "unit": item.get("unit", "") or "", - "english": item.get("english", "") or "", - "test_condition": item.get("test_condition", "") or "", - "comments": item.get("comments", "") or "", - "source_text": item.get("source_text", "") or "", - "source_page": page_label, - "chunk_type": chunk_type, - "source_text": item.get("source_text", "") or "", - }) +def convert_to_dataframe(data: Dict[str, Any]) -> pd.DataFrame: + if not data: + return pd.DataFrame() + mat_name = data.get("material_name", "") or "" + mat_abbr = data.get("material_abbreviation", "") or "" + if not mat_abbr: + mat_abbr = make_abbreviation(mat_name) - all_rows = _fuzzy_dedup(all_rows) - df = pd.DataFrame(all_rows) - if not df.empty: - base_cols = [c for c in df.columns if c not in ("source_page", "chunk_type")] - df = df[base_cols + ["source_page", "chunk_type"]] + rows = [] + for item in data.get("mechanical_properties", []): + rows.append({ + "material_name": mat_name, + "material_abbreviation": mat_abbr, + "section": item.get("section", "") or "Mechanical", + "property_name": item.get("property_name", "") or "Unknown property", + "value": item.get("value", "") or "N/A", + "unit": item.get("unit", "") or "", + "english": item.get("english", "") or "", + "test_condition": item.get("test_condition", "") or "", + "comments": item.get("comments", "") or "", + }) + return pd.DataFrame(rows) - return df, doi # ───────────────────────────────────────────────────────────────────────────── -# CONSENSUS FILTER -# Only rows present in BOTH Gemini AND GPT outputs are kept. -# Match key: normalised (section, property_name, material_name) -# Value match: exact OR within 5% numeric tolerance. +# SOURCE VERIFICATION — value found in PDF first (hard gate, tolerant), +# then meaning check on property + material (only for survivors) # ───────────────────────────────────────────────────────────────────────────── -def _first_num(s: str) -> Optional[float]: - m = re.search(r"[\d.]+", str(s)) - return float(m.group()) if m else None - - -def _value_match(val_a: str, val_b: str, tol: float = CONSENSUS_VALUE_TOL) -> bool: - na, nb = _norm(val_a), _norm(val_b) - if na == nb: - return True - fa, fb = _first_num(na), _first_num(nb) - if fa is not None and fb is not None and fb != 0: - if abs(fa - fb) / abs(fb) <= tol: - return True - # Range check: if one is a range and the other a point inside it - range_m = re.match(r"([\d.]+)[^\d.]+([\d.]+)", nb) - if range_m and fa is not None: - lo, hi = float(range_m.group(1)), float(range_m.group(2)) - if lo <= fa <= hi: - return True - return False - - -def _prop_key(row: Dict) -> Tuple[str, str, str]: - return ( - _norm(row.get("section", "")), - _norm(row.get("property_name", "")), - _norm(row.get("material_name", "")), - ) - - -def _token_set(s: str) -> set: - """Split normalised string into tokens of 3+ chars.""" - return {t for t in re.split(r"[^a-z0-9]+", _norm(s)) if len(t) >= 3} - - -def _token_overlap(a: str, b: str) -> float: - """Jaccard similarity between token sets of two strings.""" - sa, sb = _token_set(a), _token_set(b) - if not sa and not sb: - return 1.0 - if not sa or not sb: - return 0.0 - return len(sa & sb) / len(sa | sb) - +_embed_model = None -def _embed_similarity(a: str, b: str) -> float: - """Cosine similarity between two strings using the loaded SentenceTransformer.""" - try: - model = _get_embed_model() - vecs = model.encode([a, b], normalize_embeddings=True) - return float(np.dot(vecs[0], vecs[1])) - except Exception: - return 0.0 +def _get_embed_model(): + global _embed_model + if _embed_model is None: + from sentence_transformers import SentenceTransformer + _embed_model = SentenceTransformer(EMBED_MODEL_NAME) + return _embed_model -def _llm_adjudicate(pairs: List[Tuple[str, str]]) -> List[bool]: - """ - Ask GPT-4o whether each (gemini_key, gpt_key) pair refers to the same - material property. Returns a list of bools, one per pair. - Falls back to False on any error. - """ - if not pairs: +def _extract_sentences(pdf_bytes: bytes) -> List[Dict[str, Any]]: + """Plain-text sentence index of the whole PDF — used only for verification, + not for extraction. No chunking/table/ranking machinery needed here.""" + sentences = [] + with fitz.open(stream=pdf_bytes, filetype="pdf") as doc: + for page_idx, page in enumerate(doc, start=1): + raw = page.get_text("text") or "" + for para in re.split(r"(?<=[.\n])\s+", raw): + para = para.strip() + if len(para) >= 8: + sentences.append({"text": para, "page": page_idx}) + return sentences + + +def _num_close(target: float, candidate: float, rel_tol: float = VALUE_REL_TOL) -> bool: + return math.isclose(target, candidate, rel_tol=rel_tol, abs_tol=1e-6) + + +def _find_value_matches(value_str: str, sentences: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """STEP 1 — hard gate. Find every sentence containing a number within tolerance + of the target value. Tolerant, not exact-substring, so '45.3' vs '45.30' or + minor rounding differences still count as the same number.""" + val_clean = re.sub(r"^[~><=\u2248\u2265\u2264\u00b1\s]+", "", str(value_str)).strip() + m = re.search(r"[\d.]+", val_clean) + if not m: return [] - lines = "\n".join( - f"{i+1}. Gemini: \"{a}\" | GPT: \"{b}\"" - for i, (a, b) in enumerate(pairs) - ) - prompt = ( - "You are a materials science expert.\n" - "For each numbered pair below, decide if both sides refer to the " - "SAME material property (accounting for abbreviations, synonyms, " - "different units appended, or different word order).\n" - "Reply ONLY with a JSON array of booleans, one per pair, in order.\n" - "Example for 3 pairs: [true, false, true]\n\n" - f"{lines}" - ) - try: - headers = { - "Authorization": f"Bearer {OPENAI_API_KEY}", - "Content-Type": "application/json", - } - body = { - "model": GPT_MODEL, - "temperature": 0, - "max_tokens": max(len(pairs) * 15 + 100, 500), - - "messages": [{"role": "user", "content": prompt}], - } - resp = requests.post(GPT_API_URL, headers=headers, json=body, timeout=30) - if not resp.ok: - return [False] * len(pairs) - content = resp.json()["choices"][0]["message"]["content"].strip() - content = re.sub(r"^```(?:json)?\s*", "", content) - content = re.sub(r"\s*```$", "", content) - result = json.loads(content) - if isinstance(result, list) and len(result) == len(pairs): - return [bool(v) for v in result] - except Exception as e: - log.warning(f"LLM adjudication failed: {e}") - return [False] * len(pairs) - - -def _key_str(row: Dict) -> str: - """Human-readable key string for embedding/LLM — not normalised.""" - return f"{row.get('section','')} {row.get('property_name','')} {row.get('material_name','')}".strip() - - -def _cascade_match(gem_row: Dict, gpt_rows: List[Dict]) -> Optional[Dict]: - """ - Try to find a matching GPT row for a Gemini row using 4-layer cascade: - 1. Exact normalised key match - 2. Token overlap (Jaccard ≥ 0.5) - 3. Embedding cosine similarity ≥ 0.85 (clear match) - or flag as ambiguous if 0.65–0.85 (needs LLM) - 4. LLM adjudication for ambiguous pairs (called in batch outside) - Returns the best matching GPT row or None. - Also tags the row with which layer matched it. - """ - gem_key = _prop_key(gem_row) - gem_str = _key_str(gem_row) - - # Layer 1 — exact - for gpt_row in gpt_rows: - if _prop_key(gpt_row) == gem_key: - return {**gpt_row, "_match_layer": "exact"} - - # Layer 2 — token overlap - best_tok, best_tok_row = 0.0, None - for gpt_row in gpt_rows: - score = _token_overlap(gem_str, _key_str(gpt_row)) - if score > best_tok: - best_tok, best_tok_row = score, gpt_row - if best_tok >= 0.50: - return {**best_tok_row, "_match_layer": "token"} - - # Layer 3 — embedding - best_emb, best_emb_row = 0.0, None - for gpt_row in gpt_rows: - score = _embed_similarity(gem_str, _key_str(gpt_row)) - if score > best_emb: - best_emb, best_emb_row = score, gpt_row - if best_emb >= 0.85: - return {**best_emb_row, "_match_layer": "embedding"} - if best_emb >= 0.65: - # Mark as ambiguous for LLM batch — caller handles this - return {**best_emb_row, "_match_layer": "ambiguous", "_emb_score": best_emb} - - return None - -def _text_cascade_match(term: str, source_text: str, embed_threshold: float = 0.75) -> bool: - """ - Check if a term is present in source_text using cascade: - 1. Exact substring match - 2. Token overlap — any meaningful token from term appears in source text - 3. Embedding similarity between term and source text - """ - if not term or not source_text: - return False - - term_lower = term.lower().strip() - src_lower = source_text.lower().strip() - - # Layer 1 — exact substring - if term_lower in src_lower: - return True - - # Layer 2 — token overlap - tokens = [t for t in re.split(r"[^a-z0-9]+", _norm(term)) if len(t) >= 3] - if tokens and any(t in src_lower for t in tokens): - return True - - # Layer 3 — embedding similarity try: - score = _embed_similarity(term, source_text) - if score >= embed_threshold: - return True - except Exception: - pass - - return False - -def _source_text_passes(row: Dict) -> bool: - """Both property name and value must be verifiable in source_text.""" - src = str(row.get("source_text", "")).strip() - val = str(row.get("value", "")).strip() - prop = str(row.get("property_name", "")).strip() - if not src: - return False - # value check — exact numeric match only - val_ok = False - m = re.search(r"[\d.]+", val) - if m and m.group() in src: - val_ok = True - if not val_ok: - return False - # property name check — cascade - return _text_cascade_match(prop, src) - -def _resolve_value_disagreement( - gem_row: Dict, - gpt_row: Dict, - layer: str, - tol: float = CONSENSUS_VALUE_TOL, -) -> List[Dict]: - """ - Resolves a matched pair where values may differ. - Returns a list of 0, 1, or 2 rows to add to confirmed. - 0 rows = rejected - 1 row = single consensus row - 2 rows = both kept (condition variant or both unverified) - """ - gem_val = str(gem_row.get("value", "")).strip() - gpt_val = str(gpt_row.get("value", "")).strip() - gem_src = str(gem_row.get("source_text", "")).strip() - gpt_src = str(gpt_row.get("source_text", "")).strip() - gem_unit = str(gem_row.get("unit", "")).strip() - gpt_unit = str(gpt_row.get("unit", "")).strip() - gem_cond = _norm(gem_row.get("test_condition", "")) - gpt_cond = _norm(gpt_row.get("test_condition", "")) - - gem_null = not gem_val or gem_val == "N/A" - gpt_null = not gpt_val or gpt_val == "N/A" - - def _clean_gpt(row: Dict) -> Dict: - return {k: v for k, v in row.items() if not k.startswith("_")} - - def _make_row(base: Dict, other: Dict, needs_review: bool = False) -> Dict: - row = dict(base) - for k, v in other.items(): - if k.startswith("_"): - continue - if str(row.get(k, "")).strip() in ("", "N/A") and str(v).strip() not in ("", "N/A"): - row[k] = v - row["confirmed_by_gpt"] = True - row["match_layer"] = layer - row["needs_review"] = needs_review - return row - - # ── Step 1: Null fallback ───────────────────────────────────────────────── - if gem_null and not gpt_null: - if _source_text_passes(gpt_row): - row = _make_row(gem_row, gpt_row) - row["value"] = gpt_val - return [row] - else: - row = _make_row(gem_row, gpt_row, needs_review=True) - row["value"] = gpt_val - log.info(f"Null fallback unverified — '{gem_row.get('property_name','')}': GPT={gpt_val}") - return [row] - - if gpt_null and not gem_null: - if _source_text_passes(gem_row): - return [_make_row(gem_row, gpt_row)] - else: - row = _make_row(gem_row, gpt_row, needs_review=True) - log.info(f"Null fallback unverified — '{gem_row.get('property_name','')}': Gemini={gem_val}") - return [row] - - if gem_null and gpt_null: + target = float(m.group()) + except ValueError: return [] - # ── Step 2: Both non-null, values agree ─────────────────────────────────── - if _value_match(gem_val, gpt_val, tol): - return [_make_row(gem_row, gpt_row)] - - # ── Step 3: Values disagree — check conditions ──────────────────────────── - conditions_differ = gem_cond and gpt_cond and gem_cond != gpt_cond - if conditions_differ: - gem_out = dict(gem_row) - gem_out["confirmed_by_gpt"] = True - gem_out["match_layer"] = layer - gem_out["needs_review"] = False - - gpt_out = _clean_gpt(gpt_row) - gpt_out["confirmed_by_gpt"] = True - gpt_out["match_layer"] = layer - gpt_out["needs_review"] = False - - log.info( - f"Condition variant — '{gem_row.get('property_name','')}': " - f"Gemini={gem_val} ({gem_cond}), GPT={gpt_val} ({gpt_cond})" - ) - return [gem_out, gpt_out] - - # ── Step 4: Check units ─────────────────────────────────────────────────── - units_differ = gem_unit and gpt_unit and _norm(gem_unit) != _norm(gpt_unit) - if units_differ: - row = dict(gem_row) - row["confirmed_by_gpt"] = True - row["match_layer"] = layer - row["needs_review"] = False - log.info( - f"Unit mismatch — keeping Gemini '{gem_unit}' over GPT '{gpt_unit}' " - f"for '{gem_row.get('property_name','')}'" - ) - return [row] - - # ── Step 5: Source text check ───────────────────────────────────────────── - gem_verified = _source_text_passes(gem_row) - gpt_verified = _source_text_passes(gpt_row) - - if gem_verified and not gpt_verified: - log.info(f"Source text: Gemini verified — '{gem_row.get('property_name','')}': {gem_val}") - return [_make_row(gem_row, gpt_row)] - - if gpt_verified and not gem_verified: - row = _make_row(gem_row, gpt_row) - row["value"] = gpt_val - log.info(f"Source text: GPT verified — '{gem_row.get('property_name','')}': {gpt_val}") - return [row] - - # Neither verified — flag both for review - gem_out = dict(gem_row) - gem_out["confirmed_by_gpt"] = True - gem_out["match_layer"] = layer - gem_out["needs_review"] = True - - gpt_out = _clean_gpt(gpt_row) - gpt_out["confirmed_by_gpt"] = True - gpt_out["match_layer"] = layer - gpt_out["needs_review"] = True - - log.info( - f"Both unverified — '{gem_row.get('property_name','')}': " - f"Gemini={gem_val}, GPT={gpt_val} — flagged for review" - ) - return [gem_out, gpt_out] -def _post_consensus_unmatched( - df_gemini: pd.DataFrame, - df_gpt: pd.DataFrame, - df_consensus: pd.DataFrame, - tol: float = CONSENSUS_VALUE_TOL, -) -> pd.DataFrame: - """ - After consensus, collect unmatched rows from both LLMs. - Run source text check on each — keep if passes, drop if fails. - Mix kept rows silently into the consensus output. - """ - - - - def _already_in_consensus(row: Dict, consensus_rows: List[Dict]) -> bool: - """Check if this row is already represented in consensus output.""" - key = ( - _norm(row.get("material_name", "")), - _norm(row.get("property_name", "")), - _norm(row.get("value", "")), - ) - for cr in consensus_rows: - ck = ( - _norm(cr.get("material_name", "")), - _norm(cr.get("property_name", "")), - _norm(cr.get("value", "")), - ) - if key == ck: - return True - return False - - consensus_rows = df_consensus.to_dict(orient="records") if not df_consensus.empty else [] - gem_rows = df_gemini.to_dict(orient="records") if not df_gemini.empty else [] - gpt_rows = df_gpt.to_dict(orient="records") if not df_gpt.empty else [] - - kept: List[Dict] = [] - - # ── Unmatched Gemini rows ───────────────────────────────────────────────── - for row in gem_rows: - if _already_in_consensus(row, consensus_rows): - continue - if _source_text_passes(row): - clean = dict(row) - clean["confirmed_by_gpt"] = False - clean["match_layer"] = "single_gemini" - clean["needs_review"] = False - kept.append(clean) - else: - log.info( - f"Unmatched Gemini row dropped — source text failed: " - f"'{row.get('property_name','')}' = {row.get('value','')} " - f"[{row.get('material_name','')}]" - ) - - # ── Unmatched GPT rows ──────────────────────────────────────────────────── - for row in gpt_rows: - if _already_in_consensus(row, consensus_rows): - continue - if _source_text_passes(row): - clean = dict(row) - clean["confirmed_by_gpt"] = False - clean["match_layer"] = "single_gpt" - clean["needs_review"] = False - kept.append(clean) - else: - log.info( - f"Unmatched GPT row dropped — source text failed: " - f"'{row.get('property_name','')}' = {row.get('value','')} " - f"[{row.get('material_name','')}]" - ) - - if not kept: - return df_consensus - - df_extra = pd.DataFrame(kept) - - # Align columns — add missing columns from consensus with empty values - if not df_consensus.empty: - for col in df_consensus.columns: - if col not in df_extra.columns: - df_extra[col] = "" - df_extra = df_extra[df_consensus.columns] - - result = pd.concat([df_consensus, df_extra], ignore_index=True) - log.info( - f"Post-consensus: {len(kept)} unmatched rows passed source text check " - f"({sum(1 for r in kept if r['match_layer']=='single_gemini')} Gemini, " - f"{sum(1 for r in kept if r['match_layer']=='single_gpt')} GPT)" - ) - return result - -def consensus_filter( - df_gemini: pd.DataFrame, - df_gpt: pd.DataFrame, - tol: float = CONSENSUS_VALUE_TOL, -) -> pd.DataFrame: - if df_gemini.empty and df_gpt.empty: - return pd.DataFrame() - if df_gemini.empty: - return df_gpt.assign(confirmed_by_gpt=False, match_layer="none", needs_review=False) - if df_gpt.empty: - return df_gemini.assign(confirmed_by_gpt=False, match_layer="none", needs_review=False) - - from collections import defaultdict - - gem_rows = df_gemini.to_dict(orient="records") - gpt_rows = df_gpt.to_dict(orient="records") - - def _group_by_material(rows: List[Dict]) -> Dict[str, List[Dict]]: - groups: Dict[str, List[Dict]] = defaultdict(list) - for row in rows: - groups[_norm(row.get("material_name", ""))].append(row) - return groups - - gem_by_mat = _group_by_material(gem_rows) - gpt_by_mat = _group_by_material(gpt_rows) - - def _best_material_match(mat_key: str, other_groups: Dict[str, List[Dict]]) -> Optional[str]: - if mat_key in other_groups: - return mat_key - best_key, best_score = None, 0.0 - for other_key in other_groups: - score = _token_overlap(mat_key, other_key) - if score > best_score: - best_score, best_key = score, other_key - if best_score >= 0.3: - return best_key - for other_key in other_groups: - score = _embed_similarity(mat_key, other_key) - if score > best_score: - best_score, best_key = score, other_key - return best_key if best_score >= 0.70 else None - - confirmed: List[Dict] = [] - ambiguous: List[Tuple[int, Dict, Dict]] = [] - main_resolved: List[Dict] = [] - - for gem_mat_key, gem_mat_rows in gem_by_mat.items(): - matched_gpt_mat_key = _best_material_match(gem_mat_key, gpt_by_mat) - if matched_gpt_mat_key is None: - log.info(f"Material '{gem_mat_key}' found only in Gemini — skipping") - continue - - gpt_candidate_pool = gpt_by_mat[matched_gpt_mat_key] - log.info( - f"Material match: Gemini '{gem_mat_key}' → GPT '{matched_gpt_mat_key}' " - f"({len(gem_mat_rows)} Gemini rows, {len(gpt_candidate_pool)} GPT rows)" - ) - - for gem_row in gem_mat_rows: - candidate = _cascade_match(gem_row, gpt_candidate_pool) - if candidate is None: - continue - - layer = candidate.get("_match_layer", "") - if layer == "ambiguous": - ambiguous.append((len(confirmed), gem_row, candidate)) - confirmed.append(None) # slot placeholder — index must stay stable + matches = [] + for s in sentences: + for num_str in re.findall(r"[-+]?\d*\.?\d+", s["text"]): + try: + cand = float(num_str) + except ValueError: continue + if _num_close(target, cand): + matches.append(s) + break + return matches - resolved = _resolve_value_disagreement(gem_row, candidate, layer, tol) - main_resolved.extend(resolved) - - # ── LLM adjudication ───────────────────────────────────────────────────── - if ambiguous: - pairs = [(_key_str(g), _key_str(c)) for _, g, c in ambiguous] - verdicts = _llm_adjudicate(pairs) - for (slot_idx, gem_row, candidate), verdict in zip(ambiguous, verdicts): - if verdict: - resolved = _resolve_value_disagreement(gem_row, candidate, "llm", tol) - if resolved: - confirmed[slot_idx] = resolved[0] - if len(resolved) == 2: - main_resolved.append(resolved[1]) - else: - confirmed[slot_idx] = None - else: - confirmed[slot_idx] = None - - matched = [r for r in confirmed if r is not None] + main_resolved - if not matched: - log.warning("Consensus cascade: no rows agreed between Gemini and GPT-4o.") - return pd.DataFrame() - - df = pd.DataFrame(matched) - if "needs_review" not in df.columns: - df["needs_review"] = False - - layer_counts = df["match_layer"].value_counts().to_dict() - review_count = int(df["needs_review"].sum()) - log.info( - f"Consensus cascade: {len(df_gemini)} Gemini + {len(df_gpt)} GPT " - f"→ {len(df)} agreed {layer_counts} | {review_count} need review" - ) - return _post_consensus_unmatched(df_gemini, df_gpt, df, tol) - - -# ───────────────────────────────────────────────────────────────────────────── -# DOI RESOLUTION (merge from both LLMs + manual override) -# ───────────────────────────────────────────────────────────────────────────── - -def resolve_doi(doi_gemini: str, doi_gpt: str, doi_manual: str = "") -> str: - """Return the best DOI: manual override > Gemini > GPT.""" - for candidate in (doi_manual, doi_gemini, doi_gpt): - d = _normalise_doi(candidate) - if d: - return d - return "" - -# ───────────────────────────────────────────────────────────────────────────── -# CACHE HELPERS -# ───────────────────────────────────────────────────────────────────────────── - -def _df_to_cache_dict(df: pd.DataFrame, doi: str = "") -> Dict: - if df.empty: - return {} - row0 = df.iloc[0] - return { - "material_name": str(row0.get("material_name", "")), - "material_abbreviation": str(row0.get("material_abbreviation", "")), - #"trade_grade": str(row0.get("trade_grade", "")), - "manufacturer": str(row0.get("manufacturer", "")), - "doi": doi, - "mechanical_properties": df.drop( - columns=[c for c in FIXED_COLS if c in df.columns], errors="ignore" - ).to_dict(orient="records"), - } - -# ───────────────────────────────────────────────────────────────────────────── -# TOP-LEVEL PIPELINE ORCHESTRATOR -# ───────────────────────────────────────────────────────────────────────────── - -def run_pipeline( - pdf_bytes: bytes, - doi_override: str = "", - progress_callback: Any = None, -) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, List[Chunk], List[str], Dict]: - """ - Full dual-LLM consensus pipeline. - - Returns - ------- - df_consensus : rows agreed by both LLMs ← primary output - df_gemini : Gemini-only output - df_gpt : GPT-only output - all_chunks : all Chunk objects with scores - api_errors : list of error strings - meta : pipeline stats dict - """ - global GEMINI_MODEL, GEMINI_API_URL - if not GEMINI_MODEL: - key = os.getenv("GEMINI_API_KEY", "") - GEMINI_MODEL = get_gemini_model(key) - GEMINI_API_URL = ( - f"https://generativelanguage.googleapis.com/v1beta/" - f"models/{GEMINI_MODEL}:generateContent" - f"?key={key}" - ) - """ - Full dual-LLM consensus pipeline. - - Returns - ------- - df_consensus : rows agreed by both LLMs ← primary output - df_gemini : Gemini-only output - df_gpt : GPT-only output - all_chunks : all Chunk objects with scores - api_errors : list of error strings - meta : pipeline stats dict - """ - def _prog(msg: str, pct: float): - log.info(f"[{pct*100:.0f}%] {msg}") - if progress_callback: - progress_callback(msg, pct) - - meta: Dict[str, Any] = {} - pdf_hash = _pdf_hash(pdf_bytes) - - # ── Cache check ─────────────────────────────────────────────────────────── - _prog("Checking cache…", 0.0) - - - cached = cache_get(pdf_bytes) - if cached: - _prog("Cache hit.", 1.0) - meta["path"] = "cache" - df, doi = merge_to_dataframe([cached]) - doi = resolve_doi(doi, "", doi_override) - df["doi_url"] = _doi_url(doi) - # Return early — no embedding, no ChromaDB, no LLM calls - return df, df.copy(), df.copy(), [], [], meta - - # ── Stage 1: Extract ────────────────────────────────────────────────────── - _prog("Stage 1 — extracting tables + text…", 0.05) - all_chunks = extract_all_chunks(pdf_bytes) - meta["chunks_total"] = len(all_chunks) - meta["chunks_tables"] = sum(1 for c in all_chunks if c.chunk_type == "table") - meta["chunks_text"] = sum(1 for c in all_chunks if c.chunk_type == "text") - - if not all_chunks: - _prog("No content extracted.", 1.0) - meta["path"] = "failed" - return pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), [], ["No content extracted."], meta - - _prog(f"Stage 1 done — {meta['chunks_tables']} tables, {meta['chunks_text']} text.", 0.20) +def verify_dataframe(df: pd.DataFrame, sentences: List[Dict[str, Any]]) -> pd.DataFrame: + if df.empty or not sentences: + if not df.empty: + df = df.copy() + df["source_verified"] = False + return df - # ── Stage 2: Index into ChromaDB ───────────────────────────────────────── - _prog("Stage 2 — indexing into ChromaDB…", 0.25) + df = df.copy() + model = None try: - index_chunks_in_chroma(all_chunks, pdf_hash) + model = _get_embed_model() except Exception as e: - log.error(f"ChromaDB indexing failed: {e}") - - # ── Stage 3: Rank ALL chunks ────────────────────────────────────────────── - _prog("Stage 3 — ranking all chunks…", 0.40) - ranked_chunks = _chroma_rank_all(pdf_hash) - if not ranked_chunks: - # Fallback: use all chunks unranked - ranked_chunks = all_chunks - for c in ranked_chunks: - c.relevant = True - meta["chunks_ranked"] = len(ranked_chunks) - _prog(f"Stage 3 done — {len(ranked_chunks)} chunks ranked.", 0.50) - - # ── Stage 4: Build batches ──────────────────────────────────────────────── - _prog("Stage 4 — building batches…", 0.55) - batches = build_batches(ranked_chunks) - meta["batches"] = len(batches) - _prog(f"Stage 4 done — {len(batches)} batch(es).", 0.60) - - # ── Stage 5: Dual LLM in parallel ──────────────────────────────────────── - _prog("Stage 5 — Gemini + GPT-4o running in parallel…", 0.60) - gemini_raw, gpt_raw, api_errors = run_dual_llm_parallel(batches, progress_callback) - - if not gemini_raw and not gpt_raw: - meta["path"] = "failed" - return pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), all_chunks, api_errors, meta - - # ── Stage 6: Merge each LLM's results ──────────────────────────────────── - _prog("Stage 6 — merging results…", 0.90) - df_gemini, doi_gemini = merge_to_dataframe(gemini_raw) - df_gpt, doi_gpt = merge_to_dataframe(gpt_raw) - - doi = resolve_doi(doi_gemini, doi_gpt, doi_override) - doi_url = _doi_url(doi) - - meta["gemini_properties"] = len(df_gemini) - meta["gpt_properties"] = len(df_gpt) - meta["doi"] = doi - log.info(f"GPT raw batch count: {len(gpt_raw)}") - for i, r in enumerate(gpt_raw): - props = r.get("properties", r.get("mechanical_properties", [])) - log.info(f" GPT batch {i+1}: {len(props)} properties") - for p in props[:3]: # show first 3 properties from each batch - log.info(f" → {p.get('property_name','?')} = {p.get('value','?')} {p.get('unit','?')}") - - log.info(f"Gemini raw batch count: {len(gemini_raw)}") - for i, r in enumerate(gemini_raw): - props = r.get("properties", r.get("mechanical_properties", [])) - log.info(f" Gemini batch {i+1}: {len(props)} properties") - for p in props[:3]: # show first 3 properties from each batch - log.info(f" → {p.get('property_name','?')} = {p.get('value','?')} {p.get('unit','?')}") - # ── Stage 7: Consensus filter ───────────────────────────────────────────── - _prog("Stage 7 — consensus filtering…", 0.95) - df_consensus = consensus_filter(df_gemini, df_gpt) - meta["consensus_properties"] = len(df_consensus) - meta["path"] = "dual-llm-consensus" + st.warning(f"Embedding model unavailable, meaning-check skipped: {e}") + + verified = [] + for _, row in df.iterrows(): + prop = str(row.get("property_name", "") or "").strip() + mat = str(row.get("material_name", "") or "").strip() + val = str(row.get("value", "") or "").strip() + + # STEP 1 — does this value exist anywhere in the PDF (tolerant match)? + candidates = _find_value_matches(val, sentences) + if not candidates: + verified.append(False) + continue - # Attach doi_url to all frames - for df in (df_consensus, df_gemini, df_gpt): - if not df.empty: - df["doi_url"] = doi_url + # STEP 2 — of the sentences containing that value, does any one's + # meaning match the claimed property + material? + if model is None: + verified.append(True) # value confirmed; meaning check unavailable, don't over-penalize + continue + query = f"{prop} {mat}".strip() + try: + q_vec = model.encode([query], normalize_embeddings=True)[0] + texts = [c["text"] for c in candidates] + c_vecs = model.encode(texts, normalize_embeddings=True) + best_score = max(float(np.dot(q_vec, v)) for v in c_vecs) + except Exception: + best_score = 0.0 + verified.append(best_score >= MEANING_THRESHOLD) - if not df_consensus.empty: - cache_set(pdf_bytes, _df_to_cache_dict(df_consensus, doi)) + df["source_verified"] = verified + return df - _prog(f"Done — {len(df_consensus)} consensus properties.", 1.0) - return df_consensus, df_gemini, df_gpt, all_chunks, api_errors, meta # ───────────────────────────────────────────────────────────────────────────── # STREAMLIT UI # ───────────────────────────────────────────────────────────────────────────── -def _run_streamlit(): - import streamlit as st - - st.set_page_config( - page_title="DocToDB — Dual LLM Consensus Extractor", - page_icon="🧬", - layout="wide", - ) - st.title("🧬 DocToDB — Dual LLM Consensus Extractor") +def main(): + st.set_page_config(page_title="Gemini Extraction — Source Verified", layout="wide") + st.title("🧪 Gemini Extraction — Source Verified") st.caption( - f"Extracts material properties using **{GEMINI_MODEL}** + **{GPT_MODEL}** in parallel. " - "Only rows agreed by both LLMs are kept in the consensus output." + "Straight one-call Gemini extraction — no chunking, no ranking, no batching. " + "Every property is checked against the PDF: value found first, then meaning matched." ) - # ── Sidebar ─────────────────────────────────────────────────────────────── - with st.sidebar: - st.header("⚙️ Settings") - st.divider() - #st.markdown(f"**Gemini model:** `{GEMINI_MODEL}`") - #st.markdown(f"**GPT model:** `{GPT_MODEL}`") - #st.markdown(f"**Embedder:** `{EMBED_MODEL_NAME}`") - #st.markdown(f"**ChromaDB:** {'✅' if CHROMA_AVAILABLE else '❌ not installed'}") - #st.markdown(f"**Docling:** {'✅' if DOCLING_AVAILABLE else '❌'}") - #st.markdown(f"**Camelot:** {'✅' if CAMELOT_AVAILABLE else '❌'}") - #st.markdown(f"**OCR:** {'✅' if OCR_AVAILABLE else '❌'}") - gemini_ok = bool(GEMINI_API_KEY) - gpt_ok = bool(OPENAI_API_KEY) - st.markdown(f"**Gemini API Key:** {'✅' if gemini_ok else '❌ GEMINI_API_KEY not set'}") - st.markdown(f"**OpenAI API Key:** {'✅' if gpt_ok else '❌ OPENAI_API_KEY not set'}") - st.divider() - st.markdown(f"**Consensus tolerance:** ±{CONSENSUS_VALUE_TOL*100:.0f}%") - st.markdown("**Retrieval:** Schema-derived query, all chunks ranked") - st.divider() + pdf_file = st.file_uploader("PDF to extract", type=["pdf"], key="fv_pdf") - doi_manual = st.text_input( - "DOI override (optional)", - placeholder="10.1016/j.mat.2023.01.001", - help="Paste a DOI here to override auto-extraction. Leave blank to use the extracted DOI.", - ) - - st.divider() - if st.button("🗑 Clear JSON Cache"): - if os.path.exists(CACHE_FILE): - os.remove(CACHE_FILE) - st.success("JSON cache cleared.") - if st.button("🗑 Clear ChromaDB Collection"): - try: - col = _get_chroma_collection() - col.delete(where={"pdf_hash": {"$ne": ""}}) - st.success("ChromaDB cleared.") - except Exception as e: - st.error(f"ChromaDB clear failed: {e}") - - # ── Upload ──────────────────────────────────────────────────────────────── - uploaded = st.file_uploader("Upload PDF", type=["pdf"]) - if not uploaded: - st.info("Upload a PDF to get started.") + if not pdf_file: + st.info("Upload a PDF to begin.") return - pdf_bytes = uploaded.getvalue() - stem = uploaded.name.rsplit(".", 1)[0] + pdf_bytes = pdf_file.getvalue() + stem = pdf_file.name.rsplit(".", 1)[0] - if not gemini_ok: - st.error("GEMINI_API_KEY environment variable not set.") - return - if not gpt_ok: - st.warning("OPENAI_API_KEY not set — GPT-4o will be skipped; no consensus possible.") + sig = (pdf_file.name, pdf_file.size) + if st.session_state.get("fv_sig") != sig: + st.session_state["fv_sig"] = sig + st.session_state["fv_result"] = None if st.button("🚀 Run Extraction", type="primary", use_container_width=True): - bar = st.progress(0.0) - status = st.empty() - - def cb(msg, pct): - bar.progress(min(pct, 1.0)) - status.text(msg) - - with st.spinner("Running dual-LLM pipeline…"): - df_consensus, df_gemini, df_gpt, chunks, errors, meta = run_pipeline( - pdf_bytes, doi_override=doi_manual, progress_callback=cb - ) - - bar.progress(1.0) - status.empty() - - # ── DOI display ─────────────────────────────────────────────────────── - doi = meta.get("doi", "") - doi_url = _doi_url(doi) - if doi_url: - st.markdown( - f"**DOI:** [{doi}]({doi_url})", - help="Extracted from the PDF by both LLMs. Click to open the paper.", - ) - else: - st.info("No DOI found in this PDF.") - - # ── Metrics ─────────────────────────────────────────────────────────── - c1, c2, c3, c4, c5, c6, c7 = st.columns(7) - c1.metric("Tables", meta.get("chunks_tables", "—")) - c2.metric("Text blocks", meta.get("chunks_text", "—")) - c3.metric("Ranked", meta.get("chunks_ranked", "—")) - c4.metric("Batches", meta.get("batches", "—")) - c5.metric("Gemini rows", meta.get("gemini_properties", len(df_gemini))) - c6.metric("GPT-4o rows", meta.get("gpt_properties", len(df_gpt))) - c7.metric("✅ Consensus", meta.get("consensus_properties", len(df_consensus))) - - if errors: - with st.expander(f"⚠️ {len(errors)} API error(s)"): - for e in errors: - st.code(e) - - # ── Helper: render df with clickable DOI ────────────────────────────── - def _show_df_with_doi(df: pd.DataFrame, key: str): - """Display dataframe. DOI URL column rendered as markdown links.""" - if df.empty: - st.info("No data.") - return - # Build a display copy with clickable DOI - display_df = df.copy() - if "doi_url" in display_df.columns: - display_df = display_df.drop(columns=["doi_url"]) - - # Streamlit doesn't natively render HTML in dataframes, so we - # show the DOI link above and the table below. - doi_url_val = df["doi_url"].iloc[0] if "doi_url" in df.columns and not df.empty else "" - doi_val = _normalise_doi(doi_url_val.replace("https://doi.org/", "") if doi_url_val else "") - if doi_url_val: - st.markdown(f"🔗 **Paper DOI:** [{doi_val}]({doi_url_val})") - - st.dataframe(display_df, use_container_width=True, hide_index=True) - - # ── Tabs ────────────────────────────────────────────────────────────── - tab_consensus, tab_gemini, tab_gpt, tab_pages, tab_export = st.tabs([ - "✅ Consensus", - f"🟦 Gemini ({len(df_gemini)})", - f"🟩 GPT-4o ({len(df_gpt)})", - "🗂 By Page", - "📤 Export", - ]) - - with tab_consensus: - st.caption( - f"**{len(df_consensus)} properties** agreed by both Gemini and GPT-4o " - f"(fuzzy match ±{CONSENSUS_VALUE_TOL*100:.0f}%)" - ) - if df_consensus.empty: - st.warning( - "No consensus rows found. " - "Check the Gemini and GPT-4o tabs — both returned results but they don't overlap." - ) - else: - # Sub-tabs by chunk type - t_tbl, t_txt, t_all = st.tabs(["📋 Tables", "📄 Text", "🔢 All"]) - with t_tbl: - _show_df_with_doi(df_consensus[df_consensus.get("chunk_type", "") == "table"] if "chunk_type" in df_consensus.columns else pd.DataFrame(), "ct_tbl") - with t_txt: - _show_df_with_doi(df_consensus[df_consensus.get("chunk_type", "") == "text"] if "chunk_type" in df_consensus.columns else pd.DataFrame(), "ct_txt") - with t_all: - _show_df_with_doi(df_consensus, "ct_all") - - with tab_gemini: - st.caption(f"{len(df_gemini)} properties extracted by Gemini") - _show_df_with_doi(df_gemini, "gem") - - with tab_gpt: - st.caption(f"{len(df_gpt)} properties extracted by GPT-4o") - _show_df_with_doi(df_gpt, "gpt") - - with tab_pages: - df_src = df_consensus if not df_consensus.empty else df_gemini - if "source_page" in df_src.columns: - def _page_sort_key(label: str) -> int: - m = re.search(r"\d+", label) - return int(m.group()) if m else 9999 - for pg in sorted(df_src["source_page"].unique(), key=_page_sort_key): - pg_df = df_src[df_src["source_page"] == pg] - tbl_n = (pg_df.get("chunk_type", pd.Series()) == "table").sum() - txt_n = (pg_df.get("chunk_type", pd.Series()) == "text").sum() - with st.expander( - f"📄 {pg} — {len(pg_df)} " - f"propert{'y' if len(pg_df)==1 else 'ies'} " - f"({tbl_n} table · {txt_n} text)", - expanded=False, - ): - _show_df_with_doi(pg_df, f"page_{pg}") + with st.spinner("Extracting…"): + data = call_gemini_from_bytes(pdf_bytes) + df = convert_to_dataframe(data) + sentences = _extract_sentences(pdf_bytes) + st.session_state["fv_result"] = verify_dataframe(df, sentences) + + result = st.session_state.get("fv_result") + if result is None: + st.info("Click **Run Extraction** to start.") + return - with tab_export: - export_df = df_consensus if not df_consensus.empty else df_gemini - if "doi_url" in export_df.columns: - export_df = export_df.copy() + n_verified = int(result["source_verified"].sum()) if "source_verified" in result.columns else 0 + st.caption(f"{len(result)} properties extracted — {n_verified} source-verified.") + st.dataframe(result, use_container_width=True, hide_index=True) - col1, col2 = st.columns(2) - col1.download_button( - "⬇️ Consensus CSV", - export_df.to_csv(index=False).encode(), - f"{stem}_consensus.csv", - "text/csv", - use_container_width=True, - ) - col2.download_button( - "⬇️ Consensus JSON", - export_df.to_json(orient="records", indent=2).encode(), - f"{stem}_consensus.json", - "application/json", - use_container_width=True, - ) - st.divider() - col3, col4 = st.columns(2) - col3.download_button( - "⬇️ Gemini CSV", - df_gemini.to_csv(index=False).encode() if not df_gemini.empty else b"", - f"{stem}_gemini.csv", - "text/csv", - use_container_width=True, - disabled=df_gemini.empty, - ) - col4.download_button( - "⬇️ GPT-4o CSV", - df_gpt.to_csv(index=False).encode() if not df_gpt.empty else b"", - f"{stem}_gpt4o.csv", - "text/csv", - use_container_width=True, - disabled=df_gpt.empty, - ) - - st.divider() - with st.expander("🔎 Chunk inspector (top-30 by score)"): - show_type = st.radio("Show", ["all", "table", "text"], horizontal=True) - shown = [c for c in chunks if show_type == "all" or c.chunk_type == show_type] - shown = sorted(shown, key=lambda c: c.score, reverse=True)[:30] - for c in shown: - st.code( - f"[{c.chunk_type.upper()}] Page {c.page_num} | score={c.score:.3f}\n\n{c.text[:400]}", - language=None, - ) + if not result.empty: + st.download_button("⬇️ Download CSV", result.to_csv(index=False).encode(), + f"{stem}_gemini_verified.csv", "text/csv") -# ───────────────────────────────────────────────────────────────────────────── -# ENTRY POINT -# ───────────────────────────────────────────────────────────────────────────── if __name__ == "__main__": - import sys - - _in_streamlit = False - try: - import streamlit.runtime.scriptrunner as _sr - if _sr.get_script_run_ctx() is not None: - _in_streamlit = True - except Exception: - pass - - if _in_streamlit: - _run_streamlit() - else: - print( - "\nUsage:\n" - " streamlit run doctodb_dual.py\n" - "\nEnvironment variables required:\n" - " GEMINI_API_KEY — Google AI Studio key\n" - " OPENAI_API_KEY — OpenAI key\n" - ) \ No newline at end of file + main() \ No newline at end of file