File size: 1,691 Bytes
ea3fdd9 79ede11 ea3fdd9 79ede11 ea3fdd9 9c7f0b4 79ede11 25ffc8c 9c7f0b4 79ede11 25ffc8c 79ede11 25ffc8c ea3fdd9 72d3198 9c7f0b4 4e90c33 9c7f0b4 72d3198 79ede11 9c7f0b4 72d3198 9c7f0b4 79ede11 ea3fdd9 25ffc8c 72d3198 9c7f0b4 72d3198 9c7f0b4 72d3198 79ede11 72d3198 ea3fdd9 72d3198 | 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 | import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
MODEL_NAME = "Amey9766/llama32-1b-maintenance-classifier"
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
# Load model on CPU to avoid meta tensor issues
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
torch_dtype=torch.float32,
device_map=None
)
model.to("cpu")
def classify(text):
# Strong classification prompt to force label output
prompt = (
"You are a maintenance request classifier. "
"Your job is to output ONLY ONE WORD: urgent, routine, or cosmetic.\n\n"
f"Request: {text}\n"
"Category:"
)
inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
outputs = model.generate(
**inputs,
max_new_tokens=3,
temperature=0.0,
do_sample=False,
eos_token_id=tokenizer.eos_token_id
)
raw = tokenizer.decode(outputs[0], skip_special_tokens=True).lower()
# Extract only the part after "category:"
if "category:" in raw:
raw = raw.split("category:")[-1].strip()
# Match labels
if "urgent" in raw:
return "🔴 URGENT"
if "routine" in raw:
return "🟡 ROUTINE"
if "cosmetic" in raw:
return "🟢 COSMETIC"
return f"Unrecognized output: {raw}"
# Gradio UI
demo = gr.Interface(
fn=classify,
inputs=gr.Textbox(label="Enter maintenance request"),
outputs=gr.Textbox(label="Predicted Category"),
title="Maintenance Request Classifier",
description="Predicts whether a maintenance request is urgent, routine, or cosmetic."
)
if __name__ == "__main__":
demo.launch() |