Spaces:
Running on Zero
Running on Zero
File size: 2,323 Bytes
f456cd8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | from __future__ import annotations
import logging
import os
import threading
from typing import Any
import gradio as gr
import spaces
logger = logging.getLogger(__name__)
MODEL_NAME = os.getenv("NLI_MODEL_NAME", "cross-encoder/nli-deberta-v3-base")
NLI_PIPELINE = None
_nli_lock = threading.Lock()
def _get_nli_pipeline():
global NLI_PIPELINE
if NLI_PIPELINE is None:
with _nli_lock:
if NLI_PIPELINE is not None:
return NLI_PIPELINE
try:
from transformers import pipeline # type: ignore
NLI_PIPELINE = pipeline(
"text-classification",
model=MODEL_NAME,
device=-1,
)
logger.info("NLI loaded globally: %s", MODEL_NAME)
except Exception as exc:
NLI_PIPELINE = None
logger.warning("NLI unavailable globally (%s).", exc)
raise
return NLI_PIPELINE
@spaces.GPU
def score_pairs(pairs: list[dict[str, str]]) -> list[dict[str, Any]]:
pipeline = _get_nli_pipeline()
if not isinstance(pairs, list):
raise ValueError("Expected a list of {'text': premise, 'text_pair': hypothesis} objects.")
normalized_pairs: list[dict[str, str]] = []
for pair in pairs:
if not isinstance(pair, dict):
raise ValueError("Each pair must be an object.")
text = pair.get("text")
text_pair = pair.get("text_pair")
if not isinstance(text, str) or not isinstance(text_pair, str):
raise ValueError("Each pair must include string 'text' and 'text_pair' fields.")
normalized_pairs.append({"text": text, "text_pair": text_pair})
results = pipeline(normalized_pairs)
return [
{"label": str(item.get("label", "")), "score": float(item.get("score", 0.5))}
for item in results
]
demo = gr.Interface(
fn=score_pairs,
inputs=gr.JSON(
label="NLI pairs",
value=[{"text": "Paris is the capital of France.", "text_pair": "Paris is in France."}],
),
outputs=gr.JSON(label="Pipeline output"),
api_name="score_pairs",
title="AFVE NLI Service",
description="Programmatic NLI scoring service for LLMLens AFVE.",
)
if __name__ == "__main__":
demo.launch()
|