utilities/ner.py

import os
import requests

NER_API = "https://p9ai-ner.hf.space/api/ner"  # Replace with your space URL

def ner_remote(state):
    filename = state.get("filename")
    text = state.get("text")

    data = {
        "start_page": state.get("start_page", 1),
        "end_page": state.get("end_page", 1),
    }
    if text is not None:
        data["text"] = text
    if filename is not None:
        data["filename"] = filename

    headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACE_API_TOKEN')}"}
    path = state.get("temp_files", {}).get(filename)

    if path:
        with open(path, "rb") as f:
            files = {"file": (filename, f, "application/pdf")}
            resp = requests.post(NER_API, files=files, data=data, headers=headers)
    else:
        if "text" not in data:
            raise ValueError("ner_remote requires at least one of: file or text in state")
        resp = requests.post(NER_API, data=data, headers=headers)

    if resp.status_code != 200:
        raise RuntimeError(f"NER API failed: {resp.text}")

    state["ner"] = resp.json()
    return state