File size: 2,466 Bytes
e71206c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
### 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,
)
``` |