File size: 1,940 Bytes
1a1073f |
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 64 65 66 67 68 |
"""
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()
|