Update page_files/categorized/Backend/PDF_DataExtraction.py

page_files/categorized/Backend/PDF_DataExtraction.py

@@ -6,8 +6,8 @@ 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.
+Output: one table — Gemini's extracted properties, each with the paper DOI and
+flagged source_verified True/False, downloadable as CSV.
 
 Run:  streamlit run FinalVerdict.py
 Needs GEMINI_API_KEY in .env
@@ -36,6 +36,11 @@ EMBED_MODEL_NAME   = "allenai/scibert_scivocab_uncased"   # matches DocToDB_eval
 VALUE_REL_TOL      = 0.005   # 0.5% relative tolerance for numeric matching
 MEANING_THRESHOLD  = 0.35    # cosine similarity floor for property/material meaning check
 
+# DOIs follow a strict format (10.<registrant>/<suffix>), so regex over the raw
+# PDF text is far more reliable than an LLM guess. Used as the primary source;
+# Gemini's own DOI (from the schema) is only a fallback.
+DOI_CORE_RE = re.compile(r'10\.\d{4,9}/[^\s"<>\]\)]+', re.IGNORECASE)
+
 # ─────────────────────────────────────────────────────────────────────────────
 # SCHEMA / PROMPT
 # ─────────────────────────────────────────────────────────────────────────────
@@ -45,6 +50,7 @@ SCHEMA = {
     "properties": {
         "material_name": {"type": "STRING"},
         "material_abbreviation": {"type": "STRING"},
+        "doi": {"type": "STRING"},
         "mechanical_properties": {
             "type": "ARRAY",
             "items": {
@@ -66,8 +72,11 @@ SCHEMA = {
 
 EXTRACTION_PROMPT = (
     "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). "
+    "abbreviation, the paper's DOI, and ALL properties across categories (Mechanical, Thermal, "
+    "Electrical, Physical, Optical, Rheological, etc.). Return the properties as "
+    "'mechanical_properties' (a single list). "
+    "For the DOI, extract the Digital Object Identifier exactly as printed (e.g., "
+    "10.1021/acsami.0c01234); write '' if it is not present. "
     "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"
@@ -78,6 +87,7 @@ EXTRACTION_PROMPT = (
     "{\n"
     '  "material_name": "",\n'
     '  "material_abbreviation": "",\n'
+    '  "doi": "",\n'
     '  "mechanical_properties": [\n'
     '    {"section":"","property_name":"","value":"","unit":"","english":"","test_condition":"","comments":""}\n'
     "  ]\n"
@@ -93,6 +103,58 @@ def make_abbreviation(name: str) -> str:
     return abbr or name[:6].upper()
 
 
+# ─────────────────────────────────────────────────────────────────────────────
+# DOI EXTRACTION — regex over raw PDF text (primary), Gemini doi (fallback)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _clean_doi(doi: str) -> str:
+    """Strip common prefixes and trailing punctuation that regex can over-capture."""
+    if not doi:
+        return ""
+    doi = doi.strip()
+    doi = re.sub(r'^(?:https?://)?(?:dx\.)?doi\.org/', '', doi, flags=re.IGNORECASE)
+    doi = re.sub(r'^doi[:\s]+', '', doi, flags=re.IGNORECASE)
+    doi = doi.rstrip(' .,;:)]}>"\'')
+    return doi
+
+
+def _extract_doi_from_pdf(pdf_bytes: bytes) -> str:
+    """Find the paper DOI in the raw PDF text. Prefers an explicit doi.org URL or
+    a 'doi:' label, then falls back to the first bare 10.xxxx/... token."""
+    try:
+        with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
+            text = "\n".join((page.get_text("text") or "") for page in doc)
+    except Exception:
+        return ""
+
+    # 1) explicit doi.org URL
+    m = re.search(r'(?:https?://)?(?:dx\.)?doi\.org/(10\.\d{4,9}/[^\s"<>\]\)]+)', text, re.IGNORECASE)
+    if m:
+        return _clean_doi(m.group(1))
+
+    # 2) 'doi:' or 'DOI ' labelled
+    m = re.search(r'\bdoi[:\s]+\s*(10\.\d{4,9}/[^\s"<>\]\)]+)', text, re.IGNORECASE)
+    if m:
+        return _clean_doi(m.group(1))
+
+    # 3) any bare DOI token
+    m = DOI_CORE_RE.search(text)
+    if m:
+        return _clean_doi(m.group(0))
+
+    return ""
+
+
+def resolve_doi(pdf_bytes: bytes, gemini_data: Optional[Dict[str, Any]]) -> str:
+    """Prefer the regex-extracted DOI (format-verified). Fall back to Gemini's."""
+    doi = _extract_doi_from_pdf(pdf_bytes)
+    if doi:
+        return doi
+    if gemini_data:
+        return _clean_doi(str(gemini_data.get("doi", "") or ""))
+    return ""
+
+
 # ─────────────────────────────────────────────────────────────────────────────
 # GEMINI CALL — single call, whole PDF, no chunking
 # ─────────────────────────────────────────────────────────────────────────────
@@ -137,7 +199,7 @@ def call_gemini_from_bytes(pdf_bytes: bytes) -> Optional[Dict[str, Any]]:
         return None
 
 
-def convert_to_dataframe(data: Dict[str, Any]) -> pd.DataFrame:
+def convert_to_dataframe(data: Dict[str, Any], doi: str = "") -> pd.DataFrame:
     if not data:
         return pd.DataFrame()
     mat_name = data.get("material_name", "") or ""
@@ -148,6 +210,7 @@ def convert_to_dataframe(data: Dict[str, Any]) -> pd.DataFrame:
     rows = []
     for item in data.get("mechanical_properties", []):
         rows.append({
+            "doi": doi,
             "material_name": mat_name,
             "material_abbreviation": mat_abbr,
             "section": item.get("section", "") or "Mechanical",
@@ -271,10 +334,11 @@ def verify_dataframe(df: pd.DataFrame, sentences: List[Dict[str, Any]]) -> pd.Da
 
 def main():
     st.set_page_config(page_title="Gemini Extraction — Source Verified", layout="wide")
-    st.title("🧪 Gemini Extraction — Source Verified")
+    st.title(" Gemini Extraction — Source Verified")
     st.caption(
         "Straight one-call Gemini extraction — no chunking, no ranking, no batching. "
-        "Every property is checked against the PDF: value found first, then meaning matched."
+        "Every property is checked against the PDF: value found first, then meaning matched. "
+        "Paper DOI is pulled from the PDF text."
     )
 
     pdf_file = st.file_uploader("PDF to extract", type=["pdf"], key="fv_pdf")
@@ -291,24 +355,30 @@ def main():
         st.session_state["fv_sig"] = sig
         st.session_state["fv_result"] = None
 
-    if st.button("🚀 Run Extraction", type="primary", use_container_width=True):
+    if st.button(" Run Extraction", type="primary", use_container_width=True):
         with st.spinner("Extracting…"):
             data = call_gemini_from_bytes(pdf_bytes)
-            df = convert_to_dataframe(data)
+            doi = resolve_doi(pdf_bytes, data)
+            df = convert_to_dataframe(data, doi=doi)
             sentences = _extract_sentences(pdf_bytes)
             st.session_state["fv_result"] = verify_dataframe(df, sentences)
+            st.session_state["fv_doi"] = doi
 
     result = st.session_state.get("fv_result")
     if result is None:
         st.info("Click **Run Extraction** to start.")
         return
 
+    doi = st.session_state.get("fv_doi", "")
     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.caption(
+        f"DOI: {doi or '(not found)'}  —  "
+        f"{len(result)} properties extracted — {n_verified} source-verified."
+    )
     st.dataframe(result, use_container_width=True, hide_index=True)
 
     if not result.empty:
-        st.download_button("⬇️ Download CSV", result.to_csv(index=False).encode(),
+        st.download_button(" Download CSV", result.to_csv(index=False).encode(),
                             f"{stem}_gemini_verified.csv", "text/csv")