Spaces:
Sleeping
Sleeping
| """Hugging Face Spaces entry point for PaperTrail.""" | |
| import atexit | |
| import hashlib | |
| import json | |
| from importlib.metadata import distribution | |
| import os | |
| from pathlib import Path | |
| import spaces | |
| ROOT = Path(__file__).resolve().parent | |
| SOURCE_FILE = ROOT / "SOURCE.json" | |
| REPOSITORY = "https://github.com/Adarsh096/PaperTrail.git" | |
| def validate_source(): | |
| """Require the installed application to match the declared public source.""" | |
| source = json.loads(SOURCE_FILE.read_text(encoding="utf-8")) | |
| if (source.get("repository") != REPOSITORY | |
| or source.get("distribution") != "papertrail-rag" | |
| or source.get("version") != "0.1.0"): | |
| raise RuntimeError("SOURCE.json does not identify the expected PaperTrail release") | |
| revision = source.get("revision") | |
| if not isinstance(revision, str) or len(revision) != 40 or any(c not in "0123456789abcdef" for c in revision): | |
| raise RuntimeError("SOURCE.json contains an invalid source revision") | |
| installed = distribution("papertrail-rag") | |
| if installed.version != source["version"]: | |
| raise RuntimeError("Installed PaperTrail version does not match SOURCE.json") | |
| direct_url_text = installed.read_text("direct_url.json") | |
| if direct_url_text is None: | |
| raise RuntimeError("Installed PaperTrail package has no source provenance") | |
| direct_url = json.loads(direct_url_text) | |
| vcs = direct_url.get("vcs_info", {}) | |
| if direct_url.get("url") != REPOSITORY or vcs.get("commit_id") != revision: | |
| raise RuntimeError("Installed PaperTrail source does not match SOURCE.json") | |
| return source | |
| def run_gpu_models(payload): | |
| """Run BGE scoring and Qwen generation within one GPU lease.""" | |
| import time | |
| import torch | |
| from paper_evidence_rag.hf_zerogpu import answer_with_snapshots | |
| started = time.monotonic() | |
| result = answer_with_snapshots(payload) | |
| result["gpu_runtime"] = { | |
| "device": torch.cuda.get_device_name(0), | |
| "seconds": time.monotonic() - started, | |
| "peak_allocated_bytes": torch.cuda.max_memory_allocated(), | |
| } | |
| return result | |
| def configured_gpu_answerer(): | |
| if os.environ.get("PAPER_ENABLE_GPU_MODELS") != "true": | |
| return None | |
| from huggingface_hub import snapshot_download | |
| from paper_evidence_rag.hf_zerogpu import HFZeroGPUAnswerer, RERANKER_ID, GENERATOR_ID | |
| from paper_evidence_rag.reranking import MODEL_PINS as RERANKER_PINS | |
| from paper_evidence_rag.answering import MODEL_PINS as GENERATOR_PINS | |
| paths = [] | |
| for model_id, revision, name in [ | |
| (RERANKER_ID, RERANKER_PINS[RERANKER_ID], "bge"), | |
| (GENERATOR_ID, GENERATOR_PINS[GENERATOR_ID], "qwen"), | |
| ]: | |
| directory = Path("/tmp/paper-gpu-models") / name | |
| snapshot_download( | |
| model_id, revision=revision, local_dir=directory, token=False, | |
| allow_patterns=[ | |
| "config.json", "generation_config.json", "*.safetensors", | |
| "model.safetensors.index.json", "tokenizer.json", "tokenizer_config.json", | |
| "special_tokens_map.json", "vocab.json", "vocab.txt", "merges.txt", | |
| "sentencepiece.bpe.model", | |
| ], | |
| ) | |
| files = {} | |
| for path in directory.rglob("*"): | |
| if path.is_file() and ".cache" not in path.parts and path.name != "snapshot-manifest.json": | |
| with path.open("rb") as handle: | |
| files[path.relative_to(directory).as_posix()] = hashlib.file_digest(handle, "sha256").hexdigest() | |
| (directory / "snapshot-manifest.json").write_text(json.dumps({ | |
| "id": model_id, "revision": revision, "files": files, | |
| }), encoding="utf-8") | |
| paths.append(directory) | |
| answerer = HFZeroGPUAnswerer(*paths, run_gpu_models, max_seconds=90) | |
| answerer.prepare() | |
| return answerer | |
| def configured_path(name, default=None, *, required=False, must_exist=True): | |
| value = os.environ.get(name, default) | |
| if not value: | |
| if required: | |
| raise RuntimeError(f"Set the {name} Space variable") | |
| return None | |
| path = Path(value) | |
| path = path if path.is_absolute() else ROOT / path | |
| path = path.resolve() | |
| if must_exist and not path.exists(): | |
| raise RuntimeError(f"Configured path for {name} does not exist") | |
| return path | |
| def configured_cpu_encoder(embeddings): | |
| """Download the pinned CPU encoder and verify every frozen file hash.""" | |
| if os.environ.get("PAPER_ENABLE_CPU_ENCODER") != "true": | |
| return None | |
| from huggingface_hub import snapshot_download | |
| from paper_evidence_rag.isolated_encoder import IsolatedLocalEncoder | |
| metadata = json.loads(embeddings.read_text(encoding="utf-8"))["model"] | |
| directory = Path("/tmp/paper-embeddinggemma-2") | |
| snapshot_download( | |
| metadata["id"], revision=metadata["revision"], | |
| allow_patterns=list(metadata["file_sha256"]), local_dir=directory, token=False, | |
| ) | |
| for name, expected in metadata["file_sha256"].items(): | |
| path = (directory / name).resolve() | |
| if not path.is_relative_to(directory.resolve()): | |
| raise RuntimeError("Invalid encoder file path") | |
| with path.open("rb") as handle: | |
| if hashlib.file_digest(handle, "sha256").hexdigest() != expected: | |
| raise RuntimeError("Encoder hash mismatch") | |
| (directory / "snapshot-manifest.json").write_text(json.dumps({ | |
| "id": metadata["id"], "revision": metadata["revision"], | |
| "files": metadata["file_sha256"], | |
| }), encoding="utf-8") | |
| return IsolatedLocalEncoder(directory, metadata) | |
| def configured_services(budget_db): | |
| """Enable live requests only when both screening and answer providers exist.""" | |
| from paper_evidence_rag.hosted_guardrail import configured_gate | |
| groq_key = os.environ.get("GROQ_API_KEY", "").strip() | |
| openrouter_key = os.environ.get("OPENROUTER_API_KEY", "").strip() | |
| groq_enabled = os.environ.get("PAPER_ENABLE_GROQ", "").strip().lower() in {"1", "true", "yes"} | |
| openrouter_enabled = os.environ.get("PAPER_ENABLE_OPENROUTER", "").strip().lower() == "true" | |
| live_ready = groq_enabled and openrouter_enabled and bool(groq_key) and bool(openrouter_key) | |
| if not live_ready: | |
| return configured_gate(None, budget_db), None, False | |
| from paper_evidence_rag.cloud_pipeline import OpenRouterClient, PacedBudget, CloudAnswerer, SafetyFirstGate | |
| gate = configured_gate(groq_key, budget_db) | |
| cloud_client = OpenRouterClient( | |
| openrouter_key, | |
| PacedBudget("/tmp/paper-openrouter-usage.sqlite", daily_limit=30), | |
| ) | |
| atexit.register(cloud_client.close) | |
| return SafetyFirstGate(cloud_client, gate), CloudAnswerer(cloud_client), True | |
| def main(): | |
| global demo | |
| validate_source() | |
| artifacts = configured_path("PAPER_ARTIFACTS", "data/build", required=True) | |
| requests = configured_path("PAPER_REQUESTS", "outputs/requests.json") | |
| generated = configured_path("PAPER_GENERATED", "outputs/generated.json") | |
| embeddings = configured_path("PAPER_EMBEDDINGS", "outputs/embeddings.json") | |
| samples_dir = configured_path("PAPER_SAMPLES", "samples") | |
| budget_db = configured_path("PAPER_BUDGET_DB", "/tmp/paper-evidence-usage.sqlite", must_exist=False) | |
| gate, cloud_answerer, live_ready = configured_services(budget_db) | |
| sample_hashes = [hashlib.sha256(path.read_bytes()).hexdigest() | |
| for path in samples_dir.glob("*.pdf")] if samples_dir else [] | |
| from paper_evidence_rag.app_service import DemoService | |
| from paper_evidence_rag.gradio_app import build_app, launch_options | |
| from paper_evidence_rag.bucket_admission import MountedBucketAdmission | |
| admission = MountedBucketAdmission( | |
| "/data", local_lock_path="/tmp/papertrail-admission.lock", | |
| query_daily=8, visitor_query_daily=4, | |
| upload_daily=100, visitor_upload_daily=2, | |
| ) | |
| service = DemoService( | |
| artifacts, requests, generated, embeddings, gate=gate, | |
| encoder=configured_cpu_encoder(embeddings), | |
| answerer=configured_gpu_answerer(), cloud_answerer=cloud_answerer, | |
| public_pdf_hashes=sample_hashes, | |
| ) | |
| samples = {path.stem: path for path in sorted(samples_dir.glob("*.pdf"))} if samples_dir else {} | |
| if samples_dir and not samples: | |
| service.close() | |
| raise RuntimeError("PAPER_SAMPLES contains no PDF files") | |
| demo = build_app(service, samples, admission=admission, public_mode=True) | |
| import gradio as gr | |
| if not live_ready: | |
| with demo: | |
| gr.Markdown( | |
| "Live answers are temporarily unavailable. Saved examples and paper browsing remain available." | |
| ) | |
| atexit.register(service.close) | |
| options = launch_options(service, samples) | |
| options["blocked_paths"].extend([ | |
| "/data", "/tmp/paper-gpu-models", "/tmp/paper-embeddinggemma-2", | |
| "/tmp/paper-openrouter-usage.sqlite", "/tmp/paper-evidence-usage.sqlite", | |
| ]) | |
| options.pop("server_name", None) | |
| options.pop("share", None) | |
| demo.launch( | |
| server_name="0.0.0.0", server_port=int(os.environ.get("PORT", "7860")), | |
| share=False, **options, | |
| ) | |
| if __name__ == "__main__": | |
| main() | |