mARTin-B78
dd44a2c2b1
Self-hosted Docker Compose stack for DGX Spark (ARM64/aarch64, CPU-only): - SearXNG on port 8889 with JSON API enabled, rate limiting off - Firecrawl (built from local source) on port 3002, no API key required - HHEM API (FastAPI + vectara/hallucination_evaluation_model) on port 8881 - Portainer stack YAML using dgx_net external network - build-images.sh to pre-build local images before Portainer deploy - HF model cache bind-mounted from /home/sparky/LLMs/huggingface Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
77 lines
2.0 KiB
Python
77 lines
2.0 KiB
Python
from contextlib import asynccontextmanager
|
|
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel
|
|
import torch
|
|
from transformers import AutoTokenizer, AutoModelForSequenceClassification
|
|
import logging
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
MODEL_ID = "vectara/hallucination_evaluation_model"
|
|
|
|
tokenizer = None
|
|
model = None
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
global tokenizer, model
|
|
logger.info(f"Loading model {MODEL_ID} ...")
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
|
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
|
|
model.eval()
|
|
logger.info("Model ready.")
|
|
yield
|
|
del model
|
|
del tokenizer
|
|
|
|
|
|
app = FastAPI(title="HHEM API", version="1.0.0", lifespan=lifespan)
|
|
|
|
|
|
class ScoreRequest(BaseModel):
|
|
source: str
|
|
generated: str
|
|
|
|
|
|
class ScoreResponse(BaseModel):
|
|
score: float
|
|
label: str
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok", "model": MODEL_ID, "loaded": model is not None}
|
|
|
|
|
|
@app.post("/score", response_model=ScoreResponse)
|
|
def score(req: ScoreRequest):
|
|
if model is None or tokenizer is None:
|
|
raise HTTPException(status_code=503, detail="Model not loaded yet")
|
|
|
|
inputs = tokenizer(
|
|
req.source,
|
|
req.generated,
|
|
return_tensors="pt",
|
|
truncation=True,
|
|
max_length=512,
|
|
padding=True,
|
|
)
|
|
|
|
with torch.no_grad():
|
|
logits = model(**inputs).logits
|
|
probs = torch.softmax(logits, dim=-1)
|
|
|
|
# Resolve which class index maps to "Factually Consistent"
|
|
id2label = model.config.id2label # e.g. {0: "Hallucinated", 1: "Factually Consistent"}
|
|
factual_idx = next(
|
|
(idx for idx, name in id2label.items() if "consistent" in name.lower() or "factual" in name.lower()),
|
|
1, # fallback: assume class 1 is the positive (consistent) class
|
|
)
|
|
|
|
factual_score = probs[0][factual_idx].item()
|
|
label = "grounded" if factual_score > 0.5 else "hallucinated"
|
|
|
|
return ScoreResponse(score=round(factual_score, 4), label=label)
|