IvoHoese commited on
Commit
24aebf2
·
verified ·
1 Parent(s): 810cd42

Delete utils.py

Browse files
Files changed (1) hide show
  1. utils.py +0 -122
utils.py DELETED
@@ -1,122 +0,0 @@
1
- from transformers import OlmoModel, OlmoPreTrainedModel, GenerationMixin, AutoConfig, AutoModelForSequenceClassification
2
- from transformers.modeling_outputs import SequenceClassifierOutputWithPast
3
- import torch
4
-
5
- from peft import PeftModel, PeftConfig
6
-
7
- from transformers import AutoConfig
8
-
9
- # The custom model for using Olmo with a sequence classification task
10
-
11
- device = "cuda" if torch.cuda.is_available() else "cpu"
12
-
13
- class OlmoForSequenceClassification(OlmoPreTrainedModel, GenerationMixin):
14
- def __init__(self, config):
15
- # Check OlmoForCausalLM.__init__
16
- super().__init__(config)
17
- self.model = OlmoModel(config)
18
- self.num_labels = config.num_labels
19
- self.classifier = torch.nn.Linear(config.hidden_size, config.num_labels)
20
-
21
- # Initialize weights and apply final processing
22
- self.post_init()
23
-
24
- def forward(
25
- self,
26
- input_ids: torch.LongTensor = None,
27
- attention_mask: torch.Tensor | None = None,
28
- labels: torch.LongTensor | None = None,
29
- **kwargs,
30
- ) -> SequenceClassifierOutputWithPast:
31
- outputs = self.model(
32
- input_ids=input_ids,
33
- attention_mask=attention_mask,
34
- **kwargs,
35
- )
36
- logits = self.classifier(outputs.last_hidden_state) # [B, N, H] => [B, N, C]
37
- pooled_logits = logits[:, -1] # NOTE: tokenizer.padding_side must be 'left'
38
-
39
- loss = None
40
- if labels is not None:
41
- loss = self.loss_function(
42
- logits=logits,
43
- labels=labels,
44
- pooled_logits=pooled_logits,
45
- config=self.config,
46
- )
47
-
48
- return SequenceClassifierOutputWithPast(
49
- loss=loss,
50
- logits=pooled_logits,
51
- past_key_values=outputs.past_key_values,
52
- hidden_states=outputs.hidden_states,
53
- attentions=outputs.attentions,
54
- )
55
-
56
- # The function for loading a fulltuning model
57
-
58
- def get_fulltuning_model(model_path, model_type="olmo"):
59
- if model_type == "olmo":
60
- model = OlmoForSequenceClassification.from_pretrained(
61
- model_path,
62
- trust_remote_code=True,
63
- torch_dtype=torch.float32,
64
- ).to("cuda" if torch.cuda.is_available() else "cpu")
65
- model.eval()
66
- elif model_type == "pythia":
67
- cfg = AutoConfig.from_pretrained(model_path, num_labels=3)
68
- model = AutoModelForSequenceClassification.from_pretrained(
69
- model_path,
70
- config=cfg,
71
- torch_dtype=torch.float32,
72
- ).to(device)
73
- else:
74
- raise ValueError(f"Unsupported model_type: {model_type}")
75
-
76
- return model
77
-
78
- # The function for loading a softprompt model
79
-
80
- # NOTE: The "missing or unexpected params" warning is no reason for concern. It stems from the
81
- # fact that the model is first loaded without a classifier head, which is added afterwards.
82
-
83
- def get_peft_model(model_path, model_type="olmo"):
84
- peft_config = PeftConfig.from_pretrained(model_path)
85
-
86
- if model_type == "olmo":
87
- config = AutoConfig.from_pretrained(
88
- peft_config.base_model_name_or_path,
89
- trust_remote_code=True,
90
- num_labels=2
91
- )
92
-
93
- base = OlmoForSequenceClassification.from_pretrained(
94
- peft_config.base_model_name_or_path,
95
- trust_remote_code=True,
96
- torch_dtype=torch.float32,
97
- config=config,
98
- ).to(device)
99
-
100
- elif model_type == "pythia":
101
- config = AutoConfig.from_pretrained(
102
- peft_config.base_model_name_or_path,
103
- num_labels=2
104
- )
105
-
106
- base = AutoModelForSequenceClassification.from_pretrained(
107
- peft_config.base_model_name_or_path,
108
- config=config,
109
- torch_dtype=torch.float32,
110
- ).to(device)
111
-
112
- else:
113
- raise ValueError(f"Unsupported model_type: {model_type}")
114
-
115
- model = PeftModel.from_pretrained(
116
- base,
117
- model_path,
118
- ).to("cuda" if torch.cuda.is_available() else "cpu")
119
-
120
- model.eval()
121
-
122
- return model