| ### Usage | |
| ```python | |
| from transformers import LlamaModel, LlamaPreTrainedModel, TextClassificationPipeline | |
| from torch import nn | |
| import torch | |
| from typing import Dict | |
| class AtheneForSequenceClassification(LlamaPreTrainedModel): | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.model = LlamaModel(config) | |
| self.v_head = nn.Linear(config.hidden_size, 1, bias=False) | |
| self.CLS_ID = 128003 | |
| # Initialize weights and apply final processing | |
| self.post_init() | |
| def get_device(self): | |
| return self.model.device | |
| def forward( | |
| self, | |
| input_ids=None, | |
| past_key_values=None, | |
| attention_mask=None, | |
| position_ids=None, | |
| ): | |
| transformer_outputs = self.model( | |
| input_ids, | |
| attention_mask=attention_mask, | |
| position_ids=position_ids, | |
| output_hidden_states=True, | |
| ) | |
| hidden_states = transformer_outputs.hidden_states[-1] | |
| scores = [] | |
| rewards = self.v_head(hidden_states).squeeze(-1) | |
| bs = int(input_ids.shape[0]) | |
| for i in range(bs): | |
| c_inds = (input_ids[i] == self.CLS_ID).nonzero() | |
| c_ind = c_inds[-1].item() | |
| scores.append(rewards[i, c_ind]) | |
| scores = torch.stack(scores) | |
| return {"scores": scores} | |
| # Make a pipeline to handle pre and post-processing | |
| class AtheneRewardPipeline(TextClassificationPipeline): | |
| def preprocess(self, inputs, **tokenizer_kwargs) -> Dict[str, torch.Tensor]: | |
| return_tensors = self.framework | |
| formatted = self.tokenizer.apply_chat_template(inputs, tokenize=False) | |
| formatted = formatted + self.tokenizer.cls_token | |
| return self.tokenizer( | |
| formatted, | |
| return_tensors=return_tensors, | |
| max_length=4096, | |
| padding="longest", | |
| truncation=True, | |
| ) | |
| def postprocess(self, model_outputs, function_to_apply=None, top_k=1, _legacy=True): | |
| return model_outputs["scores"].cpu().float().item() | |
| # Initialize the model | |
| model = AtheneForSequenceClassification.from_pretrained("Nexusflow/Athene-RM-70B", torch_dtype=bfloat16) | |
| tokenizer = AutoTokenizer.from_pretrained("Nexusflow/Athene-RM-70B") | |
| # Initialize the pipeline | |
| pipe = pipeline( | |
| task="text-classification", | |
| model=self.model, | |
| tokenizer=self.tokenizer, | |
| pipeline_class=AtheneRewardPipeline, | |
| ) | |
| ``` |