Views
No views yet
1import torch
2from torch import nn
3from torch.nn import CrossEntropyLoss, MSELoss, BCEWithLogitsLoss
4from transformers import Gemma2PreTrainedModel,Gemma2Model, Cache, AutoTokenizer
5from transformers.modeling_outputs import SequenceClassifierOutputWithPast
6from typing import Optional, List, Union, Tuple
7from dataclasses import dataclass
8
9@dataclass
10class Config:
11 gemma_dir = 'wath5/kgl_lmsys_pref_classif'
12 max_length = 2000
13 batch_size = 8
14 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
15
16cfg = Config()
17
18class Gemma2ForSequenceClassificationV1(Gemma2PreTrainedModel):
19 def __init__(self, config):
20 super().__init__(config)
21 self.num_labels = config.num_labels
22 self.model = Gemma2Model(config)
23 self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
24
25 # Initialize weights and apply final processing
26 self.post_init()
27
28 def get_input_embeddings(self):
29 return self.model.embed_tokens
30
31 def set_input_embeddings(self, value):
32 self.model.embed_tokens = value
33
34 def forward(
35 self,
36 input_ids: torch.LongTensor = None,
37 attention_mask: Optional[torch.Tensor] = None,
38 position_ids: Optional[torch.LongTensor] = None,
39 past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
40 inputs_embeds: Optional[torch.FloatTensor] = None,
41 labels: Optional[torch.LongTensor] = None,
42 use_cache: Optional[bool] = None,
43 output_attentions: Optional[bool] = None,
44 output_hidden_states: Optional[bool] = None,
45 return_dict: Optional[bool] = None,
46 ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
47 r"""
48 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
49 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
50 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
51 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
52 """
53 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
54
55 transformer_outputs = self.model(
56 input_ids,
57 attention_mask=attention_mask,
58 position_ids=position_ids,
59 past_key_values=past_key_values,
60 inputs_embeds=inputs_embeds,
61 use_cache=use_cache,
62 output_attentions=output_attentions,
63 output_hidden_states=output_hidden_states,
64 return_dict=return_dict,
65 )
66 hidden_states = transformer_outputs[0]
67# logits = self.score(hidden_states)
68
69 if input_ids is not None:
70 batch_size = input_ids.shape[0]
71 else:
72 batch_size = inputs_embeds.shape[0]
73
74 if self.config.pad_token_id is None and batch_size != 1:
75 raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
76 if self.config.pad_token_id is None:
77 sequence_lengths = -1
78 else:
79 if input_ids is not None:
80 # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
81 sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
82 sequence_lengths = sequence_lengths % input_ids.shape[-1]
83 sequence_lengths = sequence_lengths.to(hidden_states.device)
84 else:
85 sequence_lengths = -1
86 hidden_states = hidden_states[
87 torch.arange(batch_size, device=hidden_states.device), sequence_lengths] # eos
88 pooled_logits = self.score(hidden_states)
89
90 return pooled_logits
91
92
93tokenizer = AutoTokenizer.from_pretrained("/kaggle/input/v7-dpo-16bit-01234-8bit-all/v7_dpo_16bit_01234_8bit_all")
94
95model = Gemma2ForSequenceClassificationV1.from_pretrained(
96 cfg.gemma_dir,
97 num_labels=3,
98 device_map=cfg.device,
99 use_cache=False,
100)
101model.config.pad_token_id = tokenizer.pad_token_id1from transformers.data.data_collator import pad_without_fast_tokenizer_warning
2
3@torch.no_grad()
4def single_prompt_inference(prompt, model, device, max_length=cfg.max_length):
5 """
6 Perform inference on a single prompt.
7
8 Args:
9 prompt (str): The input prompt for inference.
10 model (torch.nn.Module): The model used for inference.
11 device (torch.device): The device to run inference on.
12 tokenizer (Tokenizer): Tokenizer for preprocessing input text.
13 max_length (int): Maximum sequence length for tokenization.
14
15 Returns:
16 dict: Probabilities for "a_win", "b_win", and "tie".
17 """
18 # Tokenize the input prompt
19 input_ids = tokenizer(prompt, truncation=True, max_length=max_length)['input_ids']
20 input_ids.append(tokenizer.eos_token_id) # Add EOS token if needed
21
22 # Prepare inputs
23 inputs = pad_without_fast_tokenizer_warning(
24 tokenizer,
25 {"input_ids": [input_ids]}, # Wrap in a list for compatibility
26 padding="max_length",
27 pad_to_multiple_of=None,
28 max_length=max_length,
29 return_tensors="pt",
30 )
31
32 # Move inputs to the appropriate device
33 inputs = inputs.to(device)
34
35 # Run the model
36 outputs = model(**inputs)
37
38 # Get probabilities using softmax
39 proba = outputs.softmax(-1).cpu().squeeze()
40
41 return {
42 "winner_model_a": proba[0].item(),
43 "winner_model_b": proba[1].item(),
44 "tie": proba[2].item(),
45 }
46
47
48def create_rounds(query: str,
49 answer_a: str,
50 answer_b: str) -> str:
51 prompt =f"""User question:
52 \"""{query}\"""
53 Answer A:
54 \"""{answer_a}\"""
55 Answer B:
56 \"""{answer_b}\"""
57 """
58 return prompt
59
60query = "Hello, what is the height of the reassembled blind product?"
61answer_a = "Vous pouvez trouver toutes les informations techniques, y compris la hauteur du produit store remonté, directement sur la fiche produit de notre site. Cliquez sur l'onglet 'Produits' dnas la barre de navigation ou utilisez le moteur de recherche pour accéder au produit recherché. Avez vous une autre question ?"
62answer_b = "The height of the aluminum Venetian blind is 130 cm."
63prompt_direct = create_rounds(query, answer_a, answer_b)
64
65single_prompt_inference(prompt_direct, model=model, device=cfg.device)