from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model_name = "fiendfrye/mental-status-classifier-lama-3.1-8b-fine-tuned"

# Load model with proper quantization settings
model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    load_in_4bit=True,  # Required for this model
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Text to classify
text = "I'm trapped in a storm of emotions that I can't control, and it feels like no one understands the chaos inside me"

# Create complete prompt
prompt = f""Classify the text into Normal, Depression, Anxiety, and return the answer as the corresponding mental health disorder label.
text: {text}
label: ""

# Use pipeline for text generation
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device_map="auto",
)

outputs = pipe(prompt, max_new_tokens=2, do_sample=True, temperature=0.1)
print(outputs[0]["generated_text"].split("label: ")[-1].strip())

# Text to classify
text = "I'm trapped in a storm of emotions that I can't control, and it feels like no one understands the chaos inside me"

# Create prompt
prompt = f""Classify the text into Normal, Depression, Anxiety, and return the answer as the corresponding mental health disorder label.
text: {text}
label: ""

# Generate with model directly
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
    output = model.generate(
        **inputs,
        max_new_tokens=2,
        do_sample=True,
        temperature=0.1
    )

result = tokenizer.decode(output[0], skip_special_tokens=True)
print(result.split("label: ")[-1].strip())