1from typing import List, Optional, Union
2
3import torch
4import torch.nn as nn
5from transformers import LlamaPreTrainedModel, LlamaModel, PreTrainedTokenizerFast
6from transformers.modeling_outputs import SequenceClassifierOutputWithPast
7
8class INFORMForSequenceClassification(LlamaPreTrainedModel):
9 def __init__(self, config):
10 super().__init__(config)
11 self.num_labels = config.num_labels
12 self.model = LlamaModel(config)
13 self.score = nn.Sequential(
14 nn.Linear(config.hidden_size, config.hidden_size),
15 nn.ReLU(),
16 nn.Linear(config.hidden_size, self.num_labels)
17 )
18 # Initialize weights and apply final processing
19 self.post_init()
20
21 def forward(
22 self,
23 input_ids: Optional[torch.LongTensor] = None,
24 attention_mask: Optional[torch.Tensor] = None,
25 position_ids: Optional[torch.LongTensor] = None,
26 past_key_values: Optional[List[torch.FloatTensor]] = None,
27 inputs_embeds: Optional[torch.FloatTensor] = None,
28 labels: Optional[torch.LongTensor] = None,
29 use_cache: Optional[bool] = None,
30 output_attentions: Optional[bool] = None,
31 output_hidden_states: Optional[bool] = None,
32 return_dict: Optional[bool] = None,
33 ):
34
35 transformer_outputs = self.model(
36 input_ids,
37 attention_mask=attention_mask,
38 position_ids=position_ids,
39 past_key_values=past_key_values,
40 inputs_embeds=inputs_embeds,
41 )
42 hidden_states = transformer_outputs[0]
43 logits = self.score(hidden_states)
44
45 if input_ids is not None:
46 batch_size = input_ids.shape[0]
47 else:
48 batch_size = inputs_embeds.shape[0]
49
50 if self.config.pad_token_id is None and batch_size != 1:
51 raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
52 if self.config.pad_token_id is None:
53 sequence_lengths = -1
54 else:
55 if input_ids is not None:
56 # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
57 sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
58 sequence_lengths = sequence_lengths % input_ids.shape[-1]
59 sequence_lengths = sequence_lengths.to(logits.device)
60 else:
61 sequence_lengths = -1
62
63 pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
64
65 loss = None
66 return SequenceClassifierOutputWithPast(
67 loss=loss,
68 logits=pooled_logits,
69 past_key_values=transformer_outputs.past_key_values,
70 hidden_states=transformer_outputs.hidden_states,
71 attentions=transformer_outputs.attentions,
72 )
73
74# Load model
75model_name = "infly/INF-ORM-Llama3.1-70B"
76orm = INFORMForSequenceClassification.from_pretrained(
77 model_name,
78 torch_dtype=torch.bfloat16,
79 device_map="auto",
80 attn_implementation="flash_attention_2",
81 num_labels=1,
82)
83
84# Load tokenizer
85tokenizer = PreTrainedTokenizerFast.from_pretrained(model_name)
86conv1 = [ { "content": "Tom has 20 oranges. He gave 5 oranges to his friend Lisa, then he bought 3 more oranges. Finally, he divided all the oranges equally among himself and his 4 friends. How many oranges does each person get?", "role": "user" }, { "content": "Tom has 20 oranges. He gave 5 oranges to Lisa. 20 - 5 = 15. Tom now has 15 oranges. Tom bought 3 more oranges. 15 + 3 = 18. Tom now has 18 oranges. Tom divides the 18 oranges equally among himself and his 4 friends (a total of 5 people). 18 ÷ 5 = 3.6 oranges. Each person gets 3.6 oranges.", "role": "assistant" } ]
87conv2= [ { "content": "Tom has 20 oranges. He gave 5 oranges to his friend Lisa, then he bought 3 more oranges. Finally, he divided all the oranges equally among himself and his 4 friends. How many oranges does each person get?", "role": "user" }, { "content": "Tom has 20 oranges. He gave 5 oranges to his friend Lisa. 20 - 5 = 15. Tom now has 15 oranges. Tom bought 3 more oranges. 15 + 3 = 18. Tom now has 18 oranges. Tom divides the 18 oranges equally among his 4 friends (a total of 4 people). 18 ÷ 4 = 4.5 oranges. Each person gets 4.5 oranges.", "role": "assistant" } ]
88conv1_tokenized = tokenizer.apply_chat_template(conv1, tokenize=True, return_tensors="pt").to("cuda")
89conv2_tokenized = tokenizer.apply_chat_template(conv2, tokenize=True, return_tensors="pt").to("cuda")
90
91# Inference
92with torch.no_grad():
93 score1 = orm(conv1_tokenized).logits[0][0].item()
94 score2 = orm(conv2_tokenized).logits[0][0].item()
95print(f"Score for response 1: {score1}")
96print(f"Score for response 2: {score2}")
97
98# Output:
99# Score for response 1: 4.96875
100# Score for response 2: 2.890625
101