ohs-severity-classifier

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

model_name = "stuSterfc/ohs-severity-classifier"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

labels = ["None", "Minor", "Moderate", "Severe", "Major"]

def predict_severity(narrative: str) -> dict:
    inputs = tokenizer(
        narrative,
        return_tensors="pt",
        truncation=True,
        max_length=512,
        padding=True
    )
    
    with torch.no_grad():
        outputs = model(**inputs)
    
    probs = torch.softmax(outputs.logits, dim=1).squeeze().tolist()
    
    return {
        "probabilities": {label: round(prob, 4) 
                         for label, prob in zip(labels, probs)},
        "predicted_class": labels[probs.index(max(probs))],
        "confidence": round(max(probs), 4)
    }

# Example
result = predict_severity(
    "Employee was assisting patient transfer from bed to wheelchair. "
    "Felt sharp pain in lower back. Unable to complete shift."
)
print(result)

# IMPORTANT: Check for needlestick blind spot
needle_terms = ["needle", "needlestick", "lancet", "sharps", "puncture"]
if any(term in narrative.lower() for term in needle_terms):
    print("WARNING: Needle-related incident detected. "
          "Model has known blind spot — refer for human review.")
---

## Input Format

The model was trained on concatenated narratives using `[SEP]` tokens 
as separators between the three OSHA narrative fields, with an 
organisational size prefix: