"""
Example usage of the Motivational Interviewing BERT classifier
"""

from transformers import BertTokenizer, BertForSequenceClassification
import torch

def predict_talk_type(text, model, tokenizer):
    """Predict the talk type for a given text"""
    inputs = tokenizer(
        text, 
        return_tensors="pt", 
        padding=True, 
        truncation=True, 
        max_length=128
    )
    
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=1)
        pred = torch.argmax(probs, dim=1)
    
    label = model.config.id2label[pred.item()]
    confidence = probs[0][pred].item()
    
    return {
        'label': label,
        'confidence': confidence,
        'all_probs': {
            model.config.id2label[i]: probs[0][i].item() 
            for i in range(len(probs[0]))
        }
    }

def main():
    # Load model and tokenizer
    model_name = "RyanDDD/bert-motivational-interviewing"
    print(f"Loading model: {model_name}")
    
    tokenizer = BertTokenizer.from_pretrained(model_name)
    model = BertForSequenceClassification.from_pretrained(model_name)
    
    # Example texts
    examples = [
        "I really want to quit smoking for my health.",
        "I'm not sure if I can do this.",
        "Smoking helps me deal with stress.",
        "Maybe I should try cutting down.",
        "I've been thinking about quitting.",
        "I like smoking, it's part of who I am."
    ]
    
    print("\nPredictions:\n" + "="*60)
    
    for text in examples:
        result = predict_talk_type(text, model, tokenizer)
        
        print(f"\nText: {text}")
        print(f"Type: {result['label']} ({result['confidence']:.1%} confidence)")
        print(f"All probabilities:")
        for label, prob in result['all_probs'].items():
            print(f"  {label:8s}: {prob:.1%}")
    
    print("\n" + "="*60)

if __name__ == "__main__":
    main()