Upload 4 files

backend/llm_utils.py

@@ -0,0 +1,34 @@
+
+"""Utilities for loading the ZeroSearch simulation model and performing simulated searches."""
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import functools
+
+MODEL_NAME = "sunhaonlp/SearchSimulation_14B"
+
+@functools.lru_cache(maxsize=1)
+def _load_search_pipe():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        trust_remote_code=True,
+        device_map="auto"
+    )
+    return pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=512,
+        do_sample=False,
+        temperature=0.0,
+    )
+
+def simulate_search(query: str, k: int = 5):
+    """Generate *k* synthetic documents for *query*."""
+    pipe = _load_search_pipe()
+    prompt = f"SearchSimulation:\nQuery: {query}\nDocuments:"
+    outputs = pipe(prompt, num_return_sequences=k)
+    docs = []
+    for o in outputs:
+        text = o["generated_text"]
+        docs.append(text.split("Documents:")[-1].strip())
+    return docs

backend/main.py

@@ -0,0 +1,52 @@
+
+from fastapi import FastAPI
+from pydantic import BaseModel
+from .llm_utils import simulate_search
+from .umls_linker import link_umls
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import functools
+
+ANSWER_MODEL = "sunhaonlp/SearchSimulation_14B"
+
+@functools.lru_cache(maxsize=1)
+def _load_answer_pipe():
+    tokenizer = AutoTokenizer.from_pretrained(ANSWER_MODEL)
+    model = AutoModelForCausalLM.from_pretrained(
+        ANSWER_MODEL,
+        trust_remote_code=True,
+        device_map="auto"
+    )
+    return pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=256,
+        do_sample=False,
+        temperature=0.0,
+    )
+
+class Query(BaseModel):
+    question: str
+
+app = FastAPI(
+    title="ZeroSearch Medical Q&A API",
+    description="Ask clinical questions; get answers with UMLS links, no external search APIs.",
+    version="0.1.0",
+)
+
+@app.post("/ask")
+def ask(query: Query):
+    docs = simulate_search(query.question, k=5)
+    context = "\n\n".join(docs)
+    prompt = (
+        "Answer the medical question strictly based on the provided context.\n\n"
+        f"Context:\n{context}\n\n"
+        f"Question: {query.question}\nAnswer:"
+    )
+    answer_pipe = _load_answer_pipe()
+    answer = (
+        answer_pipe(prompt, num_return_sequences=1)[0]["generated_text"]
+        .split("Answer:")[-1].strip()
+    )
+    umls = link_umls(answer)
+    return {"answer": answer, "docs": docs, "umls": umls}

backend/umls_linker.py

@@ -0,0 +1,19 @@
+
+"""Simple UMLS linker using SciSpacy."""
+import spacy
+from scispacy.linking import UmlsEntityLinker
+
+nlp = spacy.load("en_core_sci_lg")
+linker = UmlsEntityLinker(resolve_abbreviations=True, disambiguate=True)
+nlp.add_pipe(linker)
+
+def link_umls(text: str):
+    doc = nlp(text)
+    results = []
+    for ent in doc.ents:
+        for cui, score in ent._.kb_ents:
+            results.append(
+                {"text": ent.text, "cui": cui, "score": score}
+            )
+            break  # take top candidate
+    return results