osha-tiny-llama

🦺 OSHA TinyLlama Fine-Tuned Model

Fine-tuned TinyLlama-1.1B-Chat model for Occupational Safety and Health (OSHA) risk assessment and hazard extraction from incident descriptions.

---

📦 Model Details

• Base Model: TinyLlama/TinyLlama-1.1B-Chat-v1.0
• Fine-tuned On: OSHA incident reports
• Task: JSON extraction and risk assessment
• Training Method: QLoRA (4-bit fine-tuning)
• Tokenizer: Auto-loaded from this repo

---

🚀 Intended Use

This model is designed to:

• Extract hazards, cause of accident, injury severity, and occupation from free-text incident descriptions.
• Assist safety professionals with structured risk assessments from unstructured incident logs.

---

🛠️ Usage Example

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type='nf4'
)

model = AutoModelForCausalLM.from_pretrained("your-username/osha-tiny-llama", device_map="auto", quantization_config=bnb_config)
tokenizer = AutoTokenizer.from_pretrained("your-username/osha-tiny-llama")
tokenizer.pad_token = tokenizer.eos_token

def generate_response(prompt, max_new_tokens=512):
    inputs = tokenizer(prompt, return_tensors='pt').to('cuda')
    outputs = model.generate(input_ids=inputs['input_ids'], max_new_tokens=max_new_tokens, pad_token_id=tokenizer.pad_token_id)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

prompt = """Incident Description:
A worker fell from a ladder while painting a wall.

Please extract the following in JSON format:
- Hazards
- Cause of Accident
- Degree of Injury (choose: High - may cause fatality, Medium - may cause hospitalized injury, Low - may cause non hospitalized injury)
- Occupation"""

print(generate_response(prompt))