Spaces:
Sleeping
Sleeping
| """Private phrase classification; inputs are not persisted by this application.""" | |
| import json | |
| import os | |
| from functools import lru_cache | |
| from pathlib import Path | |
| import gradio as gr | |
| from huggingface_hub import snapshot_download | |
| import numpy as np | |
| import torch | |
| import spaces | |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer | |
| from inference import calibrate, rank_average | |
| CONFIG = json.loads((Path(__file__).parent / 'deployment_config.json').read_text()) | |
| LABELS = CONFIG['label2id'] | |
| ID2LABEL = {int(index): label for label, index in LABELS.items()} | |
| ZERO_GPU = os.getenv('MS_ZERO_GPU') == '1' | |
| DEVICE = 'cuda' if ZERO_GPU or torch.cuda.is_available() else 'cpu' | |
| RECORDS = {record['name']: record for record in CONFIG['models']} | |
| CHOICES = CONFIG['choices'] | |
| def load_model(name): | |
| record = RECORDS[name] | |
| token = os.environ.get('HF_TOKEN') | |
| print('Loading classifier:', name, flush=True) | |
| model_path = snapshot_download( | |
| record['repo_id'], revision=record.get('revision'), token=token, max_workers=2, | |
| allow_patterns=['*.safetensors', 'config.json', 'tokenizer*.json', | |
| 'special_tokens_map.json', 'added_tokens.json', 'vocab.*', | |
| 'merges.txt', '*.model'], | |
| ) | |
| model = AutoModelForSequenceClassification.from_pretrained( | |
| model_path, local_files_only=True, trust_remote_code=False, | |
| use_safetensors=True, | |
| ).to(DEVICE).eval() | |
| if {label: int(index) for label, index in model.config.label2id.items()} != LABELS: | |
| raise RuntimeError('Model label mapping differs from the deployment configuration') | |
| tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True, trust_remote_code=False) | |
| print('Classifier ready:', name, flush=True) | |
| return model, tokenizer | |
| def probabilities(name, text): | |
| model, tokenizer = load_model(name) | |
| inputs = tokenizer(text, return_tensors='pt', truncation=True, | |
| max_length=int(RECORDS[name]['max_length'])).to(DEVICE) | |
| with torch.inference_mode(): | |
| return torch.softmax(model(**inputs).logits, dim=-1).cpu().numpy() | |
| def classify(text, choice): | |
| if not isinstance(text, str) or not text.strip(): | |
| raise gr.Error('Enter a phrase first.') | |
| if len(text) > 20000: | |
| raise gr.Error('Enter a shorter phrase (at most 20,000 characters).') | |
| if choice not in CHOICES: | |
| raise gr.Error('Select one of the available classifiers.') | |
| selected = CHOICES[choice] | |
| if selected == '__ensemble__': | |
| values = rank_average([probabilities(name, text) for name in CONFIG['model_order']]) | |
| values = calibrate(values, CONFIG['calibration']['temperature']) | |
| else: | |
| values = probabilities(selected, text) | |
| row = values[0] | |
| index = int(np.argmax(row)) | |
| scores = {ID2LABEL[i]: float(value) for i, value in enumerate(row)} | |
| return ID2LABEL[index], scores | |
| if ZERO_GPU: | |
| # ZeroGPU requires CUDA placement at startup, outside the GPU function. | |
| for model_name in RECORDS: | |
| load_model(model_name) | |
| with gr.Blocks(title='MS life-event classifier', analytics_enabled=False) as demo: | |
| gr.Markdown('# MS life-event classifier\nPaste one phrase and choose a classifier. Each prediction returns one event category.') | |
| phrase = gr.Textbox(label='Phrase', lines=4) | |
| choice = gr.Dropdown(list(CHOICES), value=list(CHOICES)[0], label='Classifier') | |
| button = gr.Button('Classify', variant='primary') | |
| label = gr.Textbox(label='Assigned category', interactive=False) | |
| scores = gr.Label(label='Top category scores', num_top_classes=5) | |
| gr.Markdown('Scores are diagnostic; they do not guarantee correctness. Long inputs are truncated to each model’s training token limit.') | |
| button.click(classify, [phrase, choice], [label, scores], api_name='predict') | |
| demo.queue(default_concurrency_limit=1, max_size=16) | |
| if __name__ == '__main__': | |
| demo.launch() | |