Lukas Korganas
Add AI editing conversation link
08ad240
Raw
History Blame Contribute Delete
5.21 kB
import gradio as gr
import spaces
import json
import hashlib
import logging
from functools import lru_cache
from transformers import pipeline
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
_pipeline_cache = {"hash": None, "pipe": None, "config": None}
def _hash_config(cfg: dict) -> str:
return hashlib.sha256(json.dumps(cfg, sort_keys=True).encode()).hexdigest()[:16]
@lru_cache(maxsize=3)
def _load_pipeline(hash_key: str, cfg_json: str):
cfg = json.loads(cfg_json)
logger.info(f"Loading pipeline: {cfg}")
pipe = pipeline(**cfg)
logger.info("Loaded.")
return pipe
def get_pipe(cfg: dict):
h = _hash_config(cfg)
if _pipeline_cache["hash"] == h and _pipeline_cache["pipe"] is not None:
return _pipeline_cache["pipe"], False
pipe = _load_pipeline(h, json.dumps(cfg, sort_keys=True))
_pipeline_cache.update({"hash": h, "pipe": pipe, "config": cfg})
return pipe, True
@spaces.GPU(duration=40)
def inference(pipeline_config, inputs, inference_kwargs=None):
if inference_kwargs is None:
inference_kwargs = {}
try:
pconf = json.loads(pipeline_config) if isinstance(pipeline_config, str) else pipeline_config
except Exception as e:
return json.dumps({"error": f"Bad pipeline_config: {e}"})
try:
ikw = json.loads(inference_kwargs) if isinstance(inference_kwargs, str) else inference_kwargs
except Exception as e:
return json.dumps({"error": f"Bad inference_kwargs: {e}"})
try:
pipe, reloaded = get_pipe(pconf)
except Exception as e:
return json.dumps({"error": f"Pipeline load failed: {e}"})
try:
raw = json.loads(inputs) if isinstance(inputs, str) and inputs.strip().startswith(("[", "{")) else inputs
except:
raw = inputs
try:
result = pipe(raw, **ikw)
return json.dumps({
"success": True,
"reloaded": reloaded,
"result": result,
}, default=str, indent=2)
except Exception as e:
return json.dumps({"error": f"Inference failed: {e}"})
# Example configs
EX1 = {
"pipeline_config": {"task": "text-generation", "model": "HuggingFaceTB/SmolLM2-135M-Instruct"},
"inputs": "The future of AI is",
"inference_kwargs": {"max_new_tokens": 50},
}
EX2 = {
"pipeline_config": {"task": "zero-shot-classification", "model": "facebook/bart-large-mnli"},
"inputs": "This is a contract about data privacy and user rights.",
"inference_kwargs": {"candidate_labels": ["legal", "finance", "technology", "sports"]},
}
EX3 = {
"pipeline_config": {"task": "token-classification", "model": "openai/privacy-filter"},
"inputs": "My name is Alice Smith",
"inference_kwargs": {},
}
def run_example(ex):
result = inference(ex["pipeline_config"], ex["inputs"], ex["inference_kwargs"])
return ex["pipeline_config"], ex["inputs"], ex["inference_kwargs"], result
with gr.Blocks(title="Dynamic Transformers Pipeline API") as demo:
gr.Markdown(
"# 🚀 Dynamic Transformers Pipeline API\n"
"[Continue AI conversation to edit](https://huggingface.co/chat/conversation/6a751b859a0b83e84c4abd74)\n\n"
"Zero-GPU. Pass any `transformers.pipeline` config via JSON."
)
with gr.Row():
with gr.Column(scale=1):
pcfg = gr.JSON(label="pipeline_config", value=EX1["pipeline_config"])
inp = gr.Textbox(label="inputs", value=EX1["inputs"], lines=3)
ikw = gr.JSON(label="inference_kwargs", value=EX1["inference_kwargs"])
btn = gr.Button("▶️ Run Inference", variant="primary")
with gr.Column(scale=1):
out = gr.JSON(label="output")
btn.click(inference, [pcfg, inp, ikw], out)
gr.Markdown("## 📝 Examples (click to populate & run)")
with gr.Row():
with gr.Column():
gr.Markdown("**Text Generation**")
ex1_btn = gr.Button("Run: SmolLM2 text-gen", size="sm")
ex1_out = gr.JSON(label="result")
with gr.Column():
gr.Markdown("**Zero-Shot Classification**")
ex2_btn = gr.Button("Run: zero-shot", size="sm")
ex2_out = gr.JSON(label="result")
with gr.Column():
gr.Markdown("**Token Classification (Privacy Filter)**")
ex3_btn = gr.Button("Run: privacy-filter", size="sm")
ex3_out = gr.JSON(label="result")
ex1_btn.click(
fn=lambda: run_example(EX1),
inputs=None,
outputs=[pcfg, inp, ikw, ex1_out],
)
ex2_btn.click(
fn=lambda: run_example(EX2),
inputs=None,
outputs=[pcfg, inp, ikw, ex2_out],
)
ex3_btn.click(
fn=lambda: run_example(EX3),
inputs=None,
outputs=[pcfg, inp, ikw, ex3_out],
)
gr.Markdown("## API Example")
gr.Code("""from gradio_client import Client
client = Client("DoctorSlimm/dynamic-transformers-api")
print(client.predict(
pipeline_config={"task": "text-generation", "model": "HuggingFaceTB/SmolLM2-135M-Instruct"},
inputs="The future of AI is",
inference_kwargs={"max_new_tokens": 50},
api_name="/inference"
))""", language="python")
demo.launch()