Direct llm call

evaluation/ner/direct_llm_call/neuroscience_ner_openai.py

@@ -0,0 +1,229 @@
+#!/usr/bin/env python3
+"""Run neuroscience-wide NER on an uploaded file via an OpenAI model.
+
+Usage:
+    python neuroscience_ner_openai.py --file paper.pdf --model gpt-5.5
+"""
+import argparse
+import json
+import os
+import re
+import sys
+from datetime import datetime
+
+from dotenv import load_dotenv
+from openai import OpenAI
+
+load_dotenv()
+
+# Default system/extractor prompt shipped alongside this script. Override with
+# --prompt-file to test alternative prompts.
+DEFAULT_PROMPT_FILE = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)),
+    "prompts",
+    "extractor_neuroscience_ner.txt",
+)
+
+USER_PROMPT = """\
+INPUT TEXT:
+The text to process is the attached file. Treat the full extracted text of
+the attached document as the INPUT TEXT.
+
+METADATA (paper_title / doi / source_path) — populate `source_metadata` from
+this; do NOT repeat on every entity:
+{metadata_json}
+"""
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Run neuroscience-wide NER on an uploaded file via an OpenAI model."
+    )
+    parser.add_argument("--file", "-f", required=True, help="Path to the file to upload.")
+    parser.add_argument("--model", "-m", required=True, help="Model name, e.g. gpt-5.5.")
+    parser.add_argument(
+        "--prompt-file",
+        "-p",
+        default=DEFAULT_PROMPT_FILE,
+        help=f"Path to the system/extractor prompt file. Default: {DEFAULT_PROMPT_FILE}",
+    )
+    parser.add_argument(
+        "--metadata",
+        default="{}",
+        help='JSON string with paper_title / doi / source_path. Default: "{}".',
+    )
+    parser.add_argument(
+        "--output-dir",
+        "-o",
+        default=".",
+        help="Directory to write the output JSON file into. Default: current dir.",
+    )
+    parser.add_argument(
+        "--temperature",
+        "-t",
+        type=float,
+        default=None,
+        help="Sampling temperature for reproducible runs (e.g. 0). "
+        "Omitted by default; some reasoning models reject this.",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Sampling seed for reproducible runs. Omitted by default; "
+        "some reasoning models reject this.",
+    )
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="store_true",
+        help="Print step-by-step progress and streaming updates. Default: quiet.",
+    )
+    args = parser.parse_args()
+
+    def vprint(*a, **k):
+        """Print only when --verbose is set."""
+        if args.verbose:
+            print(*a, **k)
+
+    # Parse the metadata once. We both pass it to the model (for context) and
+    # stamp it authoritatively into the saved JSON below, so the output's
+    # source_metadata is correct regardless of what the model echoes.
+    try:
+        metadata = json.loads(args.metadata)
+        if not isinstance(metadata, dict):
+            raise ValueError("--metadata must be a JSON object")
+    except (json.JSONDecodeError, ValueError) as exc:
+        parser.error(f"invalid --metadata: {exc}")
+
+    # Default source_path to the input filename when not explicitly provided.
+    metadata.setdefault("source_path", args.file)
+
+    vprint(f"Loading system prompt from: {args.prompt_file}")
+    with open(args.prompt_file, "r", encoding="utf-8") as fh:
+        system_prompt = fh.read()
+
+    vprint("Initializing OpenAI client...")
+    # Generous read timeout for very long generations; fail rather than hang forever.
+    client = OpenAI(timeout=1800.0, max_retries=2)
+
+    vprint(f"Uploading file: {args.file} ...")
+    with open(args.file, "rb") as fh:
+        uploaded = client.files.create(file=fh, purpose="user_data")
+    vprint(f"  uploaded (file_id={uploaded.id})")
+
+    vprint(f"Sending request to model '{args.model}' (streaming; this may take a while)...")
+    request_input = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "input_file", "file_id": uploaded.id},
+                {
+                    "type": "input_text",
+                    "text": USER_PROMPT.format(metadata_json=json.dumps(metadata)),
+                },
+            ],
+        }
+    ]
+
+    # Only send sampling controls when explicitly requested; some reasoning
+    # models reject temperature/seed outright.
+    request_kwargs = {}
+    if args.temperature is not None:
+        request_kwargs["temperature"] = args.temperature
+    if args.seed is not None:
+        request_kwargs["seed"] = args.seed
+    if request_kwargs:
+        vprint(f"  sampling controls: {request_kwargs}")
+
+    chunks = []
+    chars = 0
+    next_report = 2000  # print a progress line every ~2000 chars
+    with client.responses.stream(
+        model=args.model,
+        instructions=system_prompt,
+        input=request_input,
+        **request_kwargs,
+    ) as stream:
+        for event in stream:
+            if event.type == "response.output_text.delta":
+                chunks.append(event.delta)
+                chars += len(event.delta)
+                if chars >= next_report:
+                    vprint(f"  ...streaming, {chars} chars received so far")
+                    next_report += 2000
+            elif event.type == "error":
+                print(f"  stream error: {event.error}")
+        final_response = stream.get_final_response()  # surfaces any terminal API error
+
+    output_text = "".join(chunks)
+    vprint(f"  response complete — {chars} chars total.")
+
+    # Detect truncation: a low entity count is often an output-token cutoff
+    # rather than the model deciding it was done.
+    status = getattr(final_response, "status", None)
+    if status == "incomplete":
+        reason = getattr(getattr(final_response, "incomplete_details", None), "reason", None)
+        print(
+            f"WARNING: response is INCOMPLETE (reason: {reason}). "
+            "Output was truncated — entity count is an undercount. "
+            "Consider raising the model's output-token limit or chunking the input."
+        )
+
+    # Build a filesystem-safe filename with timestamp and model.
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    safe_model = re.sub(r"[^A-Za-z0-9._-]+", "-", args.model)
+    out_name = f"neuroscience_ner_{safe_model}_{timestamp}.json"
+    out_path = os.path.join(args.output_dir, out_name)
+
+    # Persist parsed JSON when possible; otherwise wrap the raw text.
+    vprint("Parsing model output as JSON...")
+    try:
+        payload = json.loads(output_text)
+        n_entities = len(payload.get("entities", [])) if isinstance(payload, dict) else 0
+        n_terms = len(payload.get("key_terms", [])) if isinstance(payload, dict) else 0
+        vprint(f"  parsed OK — {n_entities} entities, {n_terms} key_terms.")
+    except json.JSONDecodeError:
+        print("  output was not valid JSON — wrapping raw text under 'raw_output'.")
+        payload = {"raw_output": output_text}
+
+    # Stamp metadata authoritatively, overriding whatever the model echoed.
+    if isinstance(payload, dict):
+        payload["source_metadata"] = metadata
+        vprint(f"  stamped source_metadata: {metadata}")
+
+        # Compute extraction statistics and stamp them into the metadata.
+        entities = payload.get("entities", []) or []
+        label_counts = {}
+        for ent in entities:
+            if isinstance(ent, dict):
+                label = ent.get("label", "Unknown")
+                label_counts[label] = label_counts.get(label, 0) + 1
+        stats = {
+            "model": args.model,
+            "prompt_file": args.prompt_file,
+            "temperature": args.temperature,
+            "seed": args.seed,
+            "response_status": status,
+            "extracted_at": timestamp,
+            "total_entities": len(entities),
+            "entities_by_label": dict(
+                sorted(label_counts.items(), key=lambda kv: kv[1], reverse=True)
+            ),
+        }
+        payload["source_metadata"]["statistics"] = stats
+        print(
+            f"Extracted {stats['total_entities']} entities "
+            f"across {len(label_counts)} labels."
+        )
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    vprint(f"Writing output to {out_path} ...")
+    with open(out_path, "w", encoding="utf-8") as fh:
+        json.dump(payload, fh, ensure_ascii=False, indent=2)
+
+    print(f"Done. Wrote output to {out_path}")
+
+
+if __name__ == "__main__":
+    sys.exit(main())