Vivek commited on
Commit
97bea24
1 Parent(s): 28327d1

adding train files

Browse files
Files changed (3) hide show
  1. model_file.py +226 -0
  2. requirements.txt +5 -0
  3. train.py +237 -0
model_file.py ADDED
@@ -0,0 +1,226 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import jax
2
+ import jax.numpy as jnp
3
+
4
+ import flax
5
+ import flax.linen as nn
6
+ from flax.core.frozen_dict import FrozenDict, unfreeze
7
+
8
+ from typing import Any, Optional, Tuple
9
+
10
+ from transformers import (
11
+ GPT2Config)
12
+
13
+ from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
14
+ from transformers.models.gpt2.modeling_flax_gpt2 import FlaxGPT2BlockCollection
15
+ from transformers.modeling_flax_outputs import FlaxBaseModelOutput
16
+ from transformers.modeling_flax_utils import FlaxPreTrainedModel
17
+ from transformers.models.gpt2.modeling_flax_gpt2 import FlaxGPT2Module
18
+
19
+ from transformers import GPT2Tokenizer
20
+ tokenizer = GPT2Tokenizer.from_pretrained("gpt2",pad_token='<|endoftext|>')
21
+
22
+ GPT2_START_DOCSTRING = r"""
23
+ This model inherits from :class:`~transformers.FlaxPreTrainedModel`. Check the superclass documentation for the
24
+ generic methods the library implements for all its model (such as downloading or saving, resizing the input
25
+ embeddings, pruning heads etc.)
26
+ This model is also a Flax Linen `flax.nn.Module
27
+ <https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html>`__ subclass. Use it as a regular Flax
28
+ Module and refer to the Flax documentation for all matter related to general usage and behavior.
29
+ Finally, this model supports inherent JAX features such as:
30
+ - `Just-In-Time (JIT) compilation <https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit>`__
31
+ - `Automatic Differentiation <https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation>`__
32
+ - `Vectorization <https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap>`__
33
+ - `Parallelization <https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap>`__
34
+ Parameters:
35
+ config (:class:`~transformers.GPT2Config`): Model configuration class with all the parameters of the model.
36
+ Initializing with a config file does not load the weights associated with the model, only the
37
+ configuration. Check out the :meth:`~transformers.FlaxPreTrainedModel.from_pretrained` method to load the
38
+ model weights.
39
+ """
40
+ GPT2_INPUTS_DOCSTRING = r"""
41
+ Args:
42
+ input_ids (:obj:`numpy.ndarray` of shape :obj:`(batch_size,input_ids_length)`):
43
+ :obj:`input_ids_length` = ``sequence_length``. Indices of input sequence tokens in the vocabulary.
44
+ Indices can be obtained using :class:`~transformers.GPT2Tokenizer`. See
45
+ :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for
46
+ details.
47
+ `What are input IDs? <../glossary.html#input-ids>`__
48
+ attention_mask (:obj:`numpy.ndarray` of shape :obj:`(batch_size, sequence_length)`, `optional`):
49
+ Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
50
+ - 1 for tokens that are **not masked**,
51
+ - 0 for tokens that are **masked**.
52
+ `What are attention masks? <../glossary.html#attention-mask>`__
53
+ position_ids (:obj:`numpy.ndarray` of shape :obj:`(batch_size, sequence_length)`, `optional`):
54
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0,
55
+ config.max_position_embeddings - 1]``.
56
+ past_key_values (:obj:`Dict[str, np.ndarray]`, `optional`, returned by ``init_cache`` or when passing previous ``past_key_values``):
57
+ Dictionary of pre-computed hidden-states (key and values in the attention blocks) that can be used for fast
58
+ auto-regressive decoding. Pre-computed key and value hidden-states are of shape `[batch_size, max_length]`.
59
+ output_attentions (:obj:`bool`, `optional`):
60
+ Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned
61
+ tensors for more detail.
62
+ output_hidden_states (:obj:`bool`, `optional`):
63
+ Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for
64
+ more detail.
65
+ return_dict (:obj:`bool`, `optional`):
66
+ Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple.
67
+ """
68
+
69
+ class FlaxGPT2PreTrainedModel(FlaxPreTrainedModel): #modify
70
+ """
71
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
72
+ models.
73
+ """
74
+
75
+ config_class = GPT2Config
76
+ base_model_prefix = "transformer"
77
+ module_class: nn.Module = None
78
+
79
+ def __init__(
80
+ self,
81
+ config: GPT2Config,
82
+ input_shape: Tuple = (1, 1),
83
+ seed: int = 0,
84
+ dtype: jnp.dtype = jnp.float32,
85
+ **kwargs,
86
+ ):
87
+ module = self.module_class(config=config, dtype=dtype, **kwargs)
88
+ super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
89
+
90
+ def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple) -> FrozenDict:
91
+ # init input tensors
92
+ input_ids = jnp.zeros(input_shape, dtype="i4")
93
+ attention_mask = jnp.ones_like(input_ids)
94
+ position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_shape)
95
+ params_rng, dropout_rng = jax.random.split(rng)
96
+ rngs = {"params": params_rng, "dropout": dropout_rng}
97
+
98
+ return self.module.init(rngs, input_ids, attention_mask, position_ids, return_dict=False)["params"]
99
+
100
+ def init_cache(self, batch_size, max_length):
101
+ r"""
102
+ Args:
103
+ batch_size (:obj:`int`):
104
+ batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
105
+ max_length (:obj:`int`):
106
+ maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
107
+ cache.
108
+ """
109
+ # init input variables to retrieve cache
110
+ input_ids = jnp.ones((batch_size, max_length))
111
+ attention_mask = jnp.ones_like(input_ids)
112
+ position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)
113
+
114
+ init_variables = self.module.init(
115
+ jax.random.PRNGKey(0), input_ids, attention_mask, position_ids, return_dict=False, init_cache=True
116
+ )
117
+ return init_variables["cache"]
118
+
119
+ @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
120
+ def __call__(
121
+ self,
122
+ input_ids,
123
+ attention_mask=None,
124
+ position_ids=None,
125
+ params: dict = None,
126
+ past_key_values: dict = None,
127
+ dropout_rng: jax.random.PRNGKey = None,
128
+ train: bool = False,
129
+ output_attentions: Optional[bool] = None,
130
+ output_hidden_states: Optional[bool] = None,
131
+ return_dict: Optional[bool] = None,
132
+ ):
133
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
134
+ output_hidden_states = (
135
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
136
+ )
137
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
138
+
139
+
140
+ if position_ids is None:
141
+ if past_key_values is not None:
142
+ raise ValueError("Make sure to provide `position_ids` when passing `past_key_values`.")
143
+
144
+ position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)
145
+
146
+ if attention_mask is None:
147
+ attention_mask = jnp.ones_like(input_ids)
148
+
149
+ # Handle any PRNG if needed
150
+ rngs = {}
151
+ if dropout_rng is not None:
152
+ rngs["dropout"] = dropout_rng
153
+
154
+ inputs = {"params": params or self.params}
155
+
156
+ # if past_key_values are passed then cache is already initialized a private flag init_cache has to be passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that it can be changed by FlaxGPT2Attention module
157
+ if past_key_values:
158
+ inputs["cache"] = past_key_values
159
+ mutable = ["cache"]
160
+ else:
161
+ mutable = False
162
+
163
+ outputs = self.module.apply(
164
+ inputs,
165
+ jnp.array(input_ids, dtype="i4"),
166
+ jnp.array(attention_mask, dtype="i4"),
167
+ jnp.array(position_ids, dtype="i4"),
168
+ not train,
169
+ False,
170
+ output_attentions,
171
+ output_hidden_states,
172
+ return_dict,
173
+ rngs=rngs,
174
+ mutable=mutable,
175
+ )
176
+
177
+ # add updated cache to model output
178
+ if past_key_values is not None and return_dict:
179
+ outputs, past_key_values = outputs
180
+ outputs["past_key_values"] = unfreeze(past_key_values["cache"])
181
+ return outputs
182
+ elif past_key_values is not None and not return_dict:
183
+ outputs, past_key_values = outputs
184
+ outputs = outputs[:1] + (unfreeze(past_key_values["cache"]),) + outputs[1:]
185
+
186
+ return outputs
187
+
188
+ class FlaxGPT2ForMultipleChoiceModule(nn.Module):
189
+ config:GPT2Config
190
+ dtype: jnp.dtype = jnp.float32
191
+ def setup(self):
192
+ self.transformer = FlaxGPT2Module(config=self.config, dtype=self.dtype)
193
+ self.dropout = nn.Dropout(rate=0.2)
194
+ self.classifier = nn.Dense(4, dtype=self.dtype)
195
+
196
+ def __call__(self,input_ids,attention_mask,position_ids,return_dict=True,deterministic=True,*args):
197
+ batch_size = input_ids.shape[0]
198
+ rng=jax.random.PRNGKey(0)
199
+ _, dropout_rng = jax.random.split(rng)
200
+ input_ids=input_ids.reshape(4*batch_size,-1)
201
+ position_ids=position_ids.reshape(4*batch_size,-1)
202
+ attention_mask=attention_mask.reshape(4*batch_size,-1)
203
+
204
+ outputs=self.transformer(input_ids, attention_mask,position_ids,return_dict=return_dict)
205
+
206
+
207
+ hidden_states = outputs[0]
208
+ hidden_states= jnp.mean(hidden_states, axis=1)
209
+
210
+
211
+
212
+ hidden_states=hidden_states.reshape(batch_size,-1) #(32,8,768)->(32,8*768)
213
+
214
+ dropout_output = self.dropout(hidden_states,deterministic=deterministic,rng=dropout_rng)
215
+
216
+
217
+
218
+ logits = self.classifier(dropout_output)
219
+ reshaped_logits = logits.reshape(-1, 4)
220
+ #(32,4)
221
+ if not return_dict:
222
+ return (reshaped_logits,) + outputs[2:]
223
+ return reshaped_logits
224
+
225
+ class FlaxGPT2ForMultipleChoice(FlaxGPT2PreTrainedModel):
226
+ module_class = FlaxGPT2ForMultipleChoiceModule
requirements.txt ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ jax
2
+ flax
3
+ transformers
4
+ Datasets
5
+ tqdm
train.py ADDED
@@ -0,0 +1,237 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import jax
2
+ print(jax.local_device_count())
3
+ import jax.numpy as jnp
4
+
5
+ import flax
6
+ import flax.linen as nn
7
+ from flax.training.common_utils import get_metrics,onehot,shard,shard_prng_key
8
+ from flax.training import train_state
9
+ from flax.metrics.tensorboard import SummaryWriter
10
+ from flax.training import checkpoints
11
+
12
+
13
+ import logging
14
+ import optax
15
+ import math
16
+ from tqdm import tqdm
17
+
18
+ from pathlib import Path
19
+ from typing import Callable
20
+ from itertools import chain
21
+ from flax.metrics import tensorboard
22
+
23
+ from datasets import load_dataset,load_metric
24
+ from transformers import GPT2Config,GPT2Tokenizer
25
+
26
+ from model_file import FlaxGPT2ForMultipleChoice
27
+
28
+ logger = logging.getLogger()
29
+ logger.setLevel(logging.INFO)
30
+
31
+ def main():
32
+
33
+
34
+ tokenizer=GPT2Tokenizer.from_pretrained('gpt2',pad_token='<|endoftext|>')
35
+
36
+ dataset=load_dataset('cosmos_qa')
37
+
38
+ def preprocess(example):
39
+ example['context&question']=example['context']+example['question']
40
+ example['first_sentence']=[example['context&question'],example['context&question'],example['context&question'],example['context&question']]
41
+ example['second_sentence']=example['answer0'],example['answer1'],example['answer2'],example['answer3']
42
+ return example
43
+
44
+ train_dataset=dataset['train'].map(preprocess)
45
+ validation_dataset=dataset['validation'].map(preprocess)
46
+ test_dataset=dataset['test'].map(preprocess)
47
+
48
+ #Remove after experiment
49
+ len_train_dataset=64
50
+ len_validation_dataset=64
51
+ len_test_dataset=64
52
+
53
+ train_dataset=train_dataset.select(range(len_train_dataset))
54
+ test_dataset=test_dataset.select(range(len_validation_dataset))
55
+ validation_dataset=validation_dataset.select(range(len_test_dataset))
56
+
57
+ #remove_cols=train_dataset.column_names
58
+
59
+ def tokenize(examples):
60
+ a=tokenizer(examples['first_sentence'],examples['second_sentence'],padding='max_length',truncation=True,max_length=256,return_tensors='jax')
61
+ a['labels']=examples['label']
62
+ return a
63
+
64
+ train_dataset=train_dataset.map(tokenize)
65
+ validation_dataset=validation_dataset.map(tokenize)
66
+ test_dataset=test_dataset.map(tokenize)
67
+
68
+ remov_col=['id', 'context', 'question', 'answer0', 'answer1', 'answer2', 'answer3', 'labels', 'context&question', 'first_sentence', 'second_sentence']
69
+
70
+ train_dataset=train_dataset.remove_columns(remov_col)
71
+ validation_dataset=validation_dataset.remove_columns(remov_col)
72
+ test_dataset=test_dataset.remove_columns(remov_col)
73
+
74
+ per_device_batch_size=4
75
+ seed=0
76
+ num_train_epochs=1
77
+ learning_rate=2e-5
78
+
79
+
80
+ total_batch_size = per_device_batch_size * jax.local_device_count()
81
+ print('The overall batch size (both for training and eval) is', total_batch_size)
82
+ num_train_steps = len(train_dataset) // total_batch_size * num_train_epochs
83
+ num_validation_steps=len(validation_dataset)//total_batch_size*num_train_epochs
84
+
85
+ learning_rate_function = optax.linear_schedule(init_value=learning_rate, end_value=0, transition_steps=num_train_steps)
86
+
87
+ class TrainState(train_state.TrainState):
88
+ logits_function:Callable=flax.struct.field(pytree_node=False)
89
+ loss_function:Callable=flax.struct.field(pytree_node=False)
90
+
91
+ def adamw(weight_decay):
92
+ return optax.adamw(learning_rate=learning_rate_function,b1=0.9,b2=0.99,eps=1e-6,weight_decay=weight_decay)
93
+
94
+ decay_path=lambda p:not any(x in p for x in ['bias','LayerNorm.weight'])
95
+
96
+ def traverse(function):
97
+ def mask(data):
98
+ flat=flax.traverse_util.flatten_dict(data)
99
+ return flax.traverse_util.unflatten_dict({k:function(k,v) for k,v in flat.items()})
100
+ return mask
101
+
102
+ gradient_transformation=optax.chain(
103
+ optax.masked(adamw(0.0),mask=traverse(lambda path,_:decay_path(path))),
104
+ optax.masked(adamw(0.01),mask=traverse(lambda path,_:not decay_path(path))))
105
+
106
+ def loss_function(logits,labels):
107
+ logits=flax.linen.log_softmax(logits)
108
+ xentropy=optax.softmax_cross_entropy(logits,onehot(labels,num_classes=4))
109
+ return jnp.mean(xentropy)
110
+
111
+ def eval_function(logits):
112
+ return logits.argmax(-1)
113
+
114
+ model = FlaxGPT2ForMultipleChoice.from_pretrained('gpt2',input_shape=(1,4,1))
115
+
116
+ state=TrainState.create(apply_fn=model.__call__,
117
+ params=model.params,
118
+ tx=gradient_transformation,
119
+ logits_function=eval_function,
120
+ loss_function=loss_function)
121
+
122
+ def train_step(state,batch,dropout_rng):
123
+ targets=batch.pop("label")
124
+ dropout_rng,new_dropout_rng=jax.random.split(dropout_rng)
125
+ def loss_function(params):
126
+ logits=state.apply_fn(**batch,params=params,dropout_rng=dropout_rng,train=True)[0]
127
+ loss=state.loss_function(logits,targets)
128
+ return loss
129
+ grad_function=jax.value_and_grad(loss_function)
130
+ loss,grad=grad_function(state.params)
131
+ grad=jax.lax.pmean(grad,"batch")
132
+ new_state=state.apply_gradients(grads=grad)
133
+ #Added.
134
+ logits=new_state.apply_fn(**batch,params=new_state.params,dropout_rng=dropout_rng,train=True)[0]
135
+ accuracy=jnp.equal(jnp.argmax(logits,axis=-1),targets)
136
+ metrics=jax.lax.pmean({"loss":loss,"learning_rate":learning_rate_function(state.step),'accuracy':accuracy},axis_name="batch")
137
+ return new_state,metrics,new_dropout_rng
138
+
139
+ parallel_train_step = jax.pmap(train_step, axis_name="batch", donate_argnums=(0,))
140
+
141
+ def eval_step(state, batch):
142
+ targets=batch.pop('label')
143
+ logits = state.apply_fn(**batch, params=state.params, train=False)
144
+ loss=state.loss_function(logits,targets)
145
+ predictions=state.logits_function(logits)
146
+ eval_accuracy=jnp.equal(predictions,targets)
147
+ #eval_acc=jnp.equal(predictions,targets)
148
+ metrics=jax.lax.pmean({"loss":loss,'accuracy':eval_accuracy},axis_name="batch")
149
+ #return state.logits_function(logits) #(8,4)
150
+ return targets,predictions,metrics
151
+
152
+ parallel_eval_step = jax.pmap(eval_step, axis_name="batch")
153
+
154
+ def glue_train_data_loader(rng,dataset,batch_size):
155
+ steps_per_epoch=len_train_dataset//batch_size
156
+ perms=jax.random.permutation(rng,len(dataset))
157
+ perms=perms[:steps_per_epoch*batch_size]
158
+ perms=perms.reshape((steps_per_epoch,batch_size))
159
+ for perm in perms:
160
+ batch=dataset[perm]
161
+ batch={k:jnp.array(v) for k,v in batch.items()}
162
+ batch=shard(batch)
163
+ yield batch
164
+
165
+ rng=jax.random.PRNGKey(seed)
166
+ dropout_rngs=jax.random.split(rng,jax.local_device_count())
167
+
168
+ def glue_eval_data_loader(dataset, batch_size):
169
+ for i in range(len_validation_dataset // batch_size):
170
+ batch = dataset[i * batch_size : (i + 1) * batch_size]
171
+ batch = {k: jnp.array(v) for k, v in batch.items()}
172
+ batch = shard(batch)
173
+
174
+ yield batch
175
+
176
+ state = flax.jax_utils.replicate(state)
177
+ #metrics_list = list_metrics()
178
+
179
+ actual_task = "mnli"
180
+ metric = load_metric('glue', "mnli")
181
+ actual_taskmetric = load_metric('glue', actual_task)
182
+
183
+ workdir='./results_tensorboard'
184
+ summary_writer = tensorboard.SummaryWriter(workdir)
185
+ #summary_writer.hparams(dict(GPT2Config()))
186
+
187
+ logger.info(f"***** Running training *****")
188
+ logger.info(f" Num examples = {len_train_dataset}")
189
+ logger.info(f" Num Epochs = {1}")
190
+ logger.info(f" Instantaneous batch size per device = {per_device_batch_size}")
191
+ logger.info(f" Total train batch size = {total_batch_size}")
192
+ logger.info(f" Total optimization steps = {num_train_steps}")
193
+
194
+ for i, epoch in enumerate(tqdm(range(1, 2), desc=f"Epoch ...", position=0, leave=True)):
195
+ rng, input_rng = jax.random.split(rng)
196
+ train_acc_metrics=[]
197
+ train_loss_metrics=[]
198
+ eval_acc_metrics=[]
199
+ eval_loss_metrics=[]
200
+ # train
201
+ with tqdm(total=len_train_dataset // total_batch_size, desc="Training...", leave=False) as progress_bar_train:
202
+ for idx,batch in enumerate(glue_train_data_loader(input_rng, train_dataset, total_batch_size)):
203
+ state, train_metric, dropout_rngs = parallel_train_step(state, batch, dropout_rngs)
204
+ train_acc_metrics.append(jax.device_get(train_metric['accuracy']).mean().item())
205
+ train_loss_metrics.append(flax.jax_utils.unreplicate(train_metric)['loss'].item())
206
+ if idx%1==0:
207
+ summary_writer.scalar('train_loss',flax.jax_utils.unreplicate(train_metric)['loss'].item(),idx)
208
+ summary_writer.scalar('train_accuracy', jax.device_get(train_metric['accuracy']).mean().item(),idx)
209
+ if idx%1==0:
210
+ params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
211
+ model.save_pretrained(
212
+ '../gpt2-common-sense-reasoning/checkpoints',
213
+ params=params,
214
+ push_to_hub=True,
215
+ commit_message=f"Saving weights of epoch {epoch} at step {idx}",)
216
+ progress_bar_train.update(1)
217
+
218
+ # evaluate
219
+ with tqdm(total=len_validation_dataset // total_batch_size, desc="Evaluating...", leave=False) as progress_bar_eval:
220
+ for idx,batch in enumerate(glue_eval_data_loader(validation_dataset, total_batch_size)):
221
+ labels,predictions,eval_metric=parallel_eval_step(state, batch)
222
+ eval_acc_metrics.append(jax.device_get(eval_metric['accuracy']).mean().item())
223
+ eval_loss_metrics.append(flax.jax_utils.unreplicate(eval_metric)['loss'].item())
224
+ progress_bar_eval.update(1)
225
+ if idx%1==0:
226
+ logger.info(f"eval_step_loss{idx}:{flax.jax_utils.unreplicate(eval_metric)['loss'].item()} eval_step_acc{idx}:{jax.device_get(eval_metric['accuracy']).mean().item()}")
227
+ summary_writer.scalar('eval_loss',flax.jax_utils.unreplicate(eval_metric)['loss'].item(),idx)
228
+ summary_writer.scalar('eval_accuracy', jax.device_get(eval_metric['accuracy']).mean().item(),idx)
229
+
230
+ #correct
231
+ logger.info(f"Epoch {epoch} done")
232
+ logger.info(f"Train loss:{jax.device_get(jnp.array(train_loss_metrics)).mean().item()} Train accuracy:{jax.device_get(jnp.array(train_acc_metrics)).mean().item()}")
233
+ logger.info(f"Eval loss:{jax.device_get(jnp.array(eval_loss_metrics)).mean().item()} Eval accuracy:{jax.device_get(jnp.array(eval_acc_metrics)).mean().item()}")
234
+ summary_writer.flush()
235
+
236
+ if __name__ == "__main__":
237
+ main()