File size: 4,044 Bytes
810cd42 | 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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | from transformers import OlmoModel, OlmoPreTrainedModel, GenerationMixin, AutoConfig, AutoModelForSequenceClassification
from transformers.modeling_outputs import SequenceClassifierOutputWithPast
import torch
from peft import PeftModel, PeftConfig
from transformers import AutoConfig
# The custom model for using Olmo with a sequence classification task
device = "cuda" if torch.cuda.is_available() else "cpu"
class OlmoForSequenceClassification(OlmoPreTrainedModel, GenerationMixin):
def __init__(self, config):
# Check OlmoForCausalLM.__init__
super().__init__(config)
self.model = OlmoModel(config)
self.num_labels = config.num_labels
self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
# Initialize weights and apply final processing
self.post_init()
def forward(
self,
input_ids: torch.LongTensor = None,
attention_mask: torch.Tensor | None = None,
labels: torch.LongTensor | None = None,
**kwargs,
) -> SequenceClassifierOutputWithPast:
outputs = self.model(
input_ids=input_ids,
attention_mask=attention_mask,
**kwargs,
)
logits = self.classifier(outputs.last_hidden_state) # [B, N, H] => [B, N, C]
pooled_logits = logits[:, -1] # NOTE: tokenizer.padding_side must be 'left'
loss = None
if labels is not None:
loss = self.loss_function(
logits=logits,
labels=labels,
pooled_logits=pooled_logits,
config=self.config,
)
return SequenceClassifierOutputWithPast(
loss=loss,
logits=pooled_logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
)
# The function for loading a fulltuning model
def get_fulltuning_model(model_path, model_type="olmo"):
if model_type == "olmo":
model = OlmoForSequenceClassification.from_pretrained(
model_path,
trust_remote_code=True,
torch_dtype=torch.float32,
).to("cuda" if torch.cuda.is_available() else "cpu")
model.eval()
elif model_type == "pythia":
cfg = AutoConfig.from_pretrained(model_path, num_labels=3)
model = AutoModelForSequenceClassification.from_pretrained(
model_path,
config=cfg,
torch_dtype=torch.float32,
).to(device)
else:
raise ValueError(f"Unsupported model_type: {model_type}")
return model
# The function for loading a softprompt model
# NOTE: The "missing or unexpected params" warning is no reason for concern. It stems from the
# fact that the model is first loaded without a classifier head, which is added afterwards.
def get_peft_model(model_path, model_type="olmo"):
peft_config = PeftConfig.from_pretrained(model_path)
if model_type == "olmo":
config = AutoConfig.from_pretrained(
peft_config.base_model_name_or_path,
trust_remote_code=True,
num_labels=2
)
base = OlmoForSequenceClassification.from_pretrained(
peft_config.base_model_name_or_path,
trust_remote_code=True,
torch_dtype=torch.float32,
config=config,
).to(device)
elif model_type == "pythia":
config = AutoConfig.from_pretrained(
peft_config.base_model_name_or_path,
num_labels=2
)
base = AutoModelForSequenceClassification.from_pretrained(
peft_config.base_model_name_or_path,
config=config,
torch_dtype=torch.float32,
).to(device)
else:
raise ValueError(f"Unsupported model_type: {model_type}")
model = PeftModel.from_pretrained(
base,
model_path,
).to("cuda" if torch.cuda.is_available() else "cpu")
model.eval()
return model |