Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3
4model_id = "weiminw/Heliumos-RM-4B"
5tokenizer = AutoTokenizer.from_pretrained(model_id,padding_side="left")
6model = AutoModelForSequenceClassification.from_pretrained(
7 model_id,
8 torch_dtype="auto",
9 device_map="auto"
10 )
11
12messages = [
13 {'role': 'user', 'content': "what is 12 * 12?"},
14 {'role': 'assistant', 'content': "144"}
15
16]
17
18text_encoded = tokenizer.apply_chat_template(messages, return_dict=True, return_tensors="pt", tokenize=True).to("cuda:0")
19
20score = model(**text_encoded)
21
22print(score.logits[0]) #
23