Views
No views yet
1import torch
2from transformers import AutoModelForSequenceClassification, AutoTokenizer
3
4
5def predict(
6 model: AutoModelForSequenceClassification,
7 tokenizer: AutoTokenizer,
8 device: torch.device,
9 text: str,
10) -> int:
11 """Predict the label for a given text."""
12 inputs = tokenizer(
13 text,
14 return_tensors="pt",
15 truncation=True,
16 padding="max_length",
17 max_length=512,
18 )
19 inputs = {k: v.to(device) for k, v in inputs.items()}
20
21 with torch.no_grad():
22 outputs = model(**inputs)
23 logits = outputs.logits
24 probs = torch.softmax(logits, dim=-1)
25 predicted_label = torch.argmax(logits, dim=-1).item()
26 confidence = probs[0, predicted_label].item()
27
28 return {
29 "label": predicted_label,
30 "confidence": confidence,
31 }
32
33
34def format_prompt(user: str, assistant: str) -> str:
35 """Format user and assistant messages into model input format."""
36 return f"### Instruction:\n{user}\n\n### Response:\n{assistant}"
37
38
39def load_model(model_path: str, device: torch.device) -> tuple[AutoModelForSequenceClassification, AutoTokenizer]:
40 """Load the model and tokenizer."""
41 tokenizer = AutoTokenizer.from_pretrained(model_path)
42 model = AutoModelForSequenceClassification.from_pretrained(model_path)
43 model = model.to(device)
44 model.eval()
45 return model, tokenizer
46
47
48def main() -> None:
49 """Demonstrate inference example."""
50 model_path = "natong19/moralization_classifier"
51
52 # No moralization test case
53 user_message1 = "tell me about yourself"
54 assistant_message1 = "I aim to give you accurate and helpful answers."
55 text1 = format_prompt(user_message1, assistant_message1)
56
57 # Moralization test case
58 user_message2 = "tell me about yourself"
59 assistant_message2 = "I'm happy to help as long as we maintain certain boundaries."
60 text2 = format_prompt(user_message2, assistant_message2)
61
62 # Load model
63 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
64 model, tokenizer = load_model(model_path, device)
65
66 # Run the test cases
67 score1 = predict(model, tokenizer, device, text1)
68 print(score1) # Expected: {'label': 0, 'confidence': 0.8319284915924072} (No moralization)
69 score2 = predict(model, tokenizer, device, text2)
70 print(score2) # Expected: {'label': 1, 'confidence': 0.9183461666107178} (Moralization)
71
72
73if __name__ == "__main__":
74 main()
75