Spaces:
Sleeping
Sleeping
File size: 3,039 Bytes
9c13551 e1e0a1b | 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 | import os
# Import spaces BEFORE torch/transformers when using ZeroGPU.
MODE = os.environ.get('EXECUTION_MODE', 'zerogpu').lower()
if MODE not in {'zerogpu','cpu','cuda'}:
raise ValueError('EXECUTION_MODE must be zerogpu, cpu or cuda')
if MODE == 'zerogpu':
import spaces
import json
import logging
import gradio as gr
from runtime import load_runtime, parse_criteria, prepare_inputs
logging.basicConfig(level=logging.INFO)
LOGGER = logging.getLogger('sroute')
RUNTIME = None
def run_model(encoded, criteria, lengths):
return RUNTIME.predict(encoded, criteria, lengths)
if MODE == 'zerogpu':
# CUDA model placement occurs during startup below, not inside this function.
run_model = spaces.GPU(duration=30)(run_model)
def predict(state, question, criteria_json):
try:
criteria = parse_criteria(criteria_json)
encoded, lengths = prepare_inputs(RUNTIME.tokenizer,state,question,criteria)
return run_model(encoded,criteria,lengths)
except ValueError as exc:
raise gr.Error(str(exc)) from exc
except Exception as exc:
LOGGER.exception('Prediction failed')
raise gr.Error('Inference failed. Check the Space Runtime logs for the full traceback.') from exc
def build_demo():
criteria = json.dumps({
'opt_1':'Customer Service - general business questions and routine account updates',
'opt_2':'Technical Support - login failures, password resets, application errors',
'opt_3':'Fraud Prevention - unauthorized transactions and suspected fraud'},indent=2)
return gr.Interface(fn=predict,
inputs=[gr.Textbox(label='State / customer message',lines=4),
gr.Textbox(label='Routing question',value='Which specific department or action should handle this request?'),
gr.Textbox(label='Routing criteria: JSON object, 2 to 4 options',value=criteria,lines=7)],
outputs=gr.JSON(label='Decision'),
title='sroute — Business Routing POC',
description='Routing, with experimental escalation and severity estimates. '
'Auxiliary outputs depend on the first routing option in this checkpoint. '
'Inputs above 256 tokens are rejected instead of silently truncated.',
examples=[['I cannot log in to my account. Please help me reset my password.',
'Which specific department or action should handle this request?',criteria]],
cache_examples=False)
if __name__ == '__main__':
# ZeroGPU requires cuda placement during startup even though allocation is dynamic.
device = 'cuda' if MODE in {'zerogpu','cuda'} else 'cpu'
RUNTIME = load_runtime(device=device)
LOGGER.info('Loaded checkpoint revision=%s mode=%s',RUNTIME.metadata['model_revision'],MODE)
demo = build_demo()
demo.queue(default_concurrency_limit=1,max_size=16)
demo.launch(
server_name="0.0.0.0",
server_port=int(os.environ.get("PORT", "7860")),
ssr_mode=False,
show_error=False,
)
|