uploaded prediction

src/prediction.py

@@ -0,0 +1,57 @@
+import jax
+import jax.numpy as jnp
+
+import flax
+import flax.linen as nn
+from flax.core.frozen_dict import FrozenDict, unfreeze
+
+from typing import Any, Optional, Tuple
+
+from transformers import (
+    GPT2Config)
+
+import transformers
+from transformers import GPT2Tokenizer
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2",pad_token='<|endoftext|>') 
+from datasets import load_dataset,load_metric
+
+from datasets import Dataset
+
+from model_file import FlaxGPT2ForMultipleChoice
+
+import logging
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+run_dataset=Dataset.from_csv('......')
+
+def preprocess(example):
+    example['context&question']=example['context']+example['question']
+    example['first_sentence']=[example['context&question'],example['context&question'],example['context&question'],example['context&question']]
+    example['second_sentence']=example['answer0'],example['answer1'],example['answer2'],example['answer3']
+    return example
+
+run_dataset=run_dataset.map(preprocess)
+
+def tokenize(examples):
+    a=tokenizer(examples['first_sentence'],examples['second_sentence'],padding='max_length',truncation=True,max_length=256,return_tensors='jax')
+    a['labels']=examples['label']
+    return a
+
+run_dataset=run_dataset.map(tokenize)
+
+remov_col=['id', 'context', 'question', 'answer0', 'answer1', 'answer2', 'answer3', 'labels', 'context&question', 'first_sentence', 'second_sentence']
+
+run_dataset=run_dataset.remove_columns(remov_col)
+
+model = FlaxGPT2ForMultipleChoice.from_pretrained("flax-community/gpt2-Cosmos",input_shape=(1,4,1))
+
+input_id=jnp.array(run_dataset['input_ids'])
+att_mask=jnp.array(run_dataset['attention_mask'])
+
+outputs=model(input_id,att_mask)
+
+final_output=jnp.argmax(outputs,axis=-1)
+
+logger.info(f"the predction of the dataset : {final_output}")

src/test.py

@@ -0,0 +1,59 @@
+import jax
+import jax.numpy as jnp
+
+import flax
+import flax.linen as nn
+from flax.core.frozen_dict import FrozenDict, unfreeze
+
+from typing import Any, Optional, Tuple
+
+from transformers import (
+    GPT2Config)
+
+import transformers
+from transformers import GPT2Tokenizer
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2",pad_token='<|endoftext|>') 
+from datasets import load_dataset,load_metric
+
+from model_file import FlaxGPT2ForMultipleChoice
+
+import logging
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+dataset=load_dataset('cosmos_qa')
+
+len_test_dataset=6963
+
+test_dataset=dataset['test'].select(range(len_test_dataset))
+
+def preprocess(example):
+    example['context&question']=example['context']+example['question']
+    example['first_sentence']=[example['context&question'],example['context&question'],example['context&question'],example['context&question']]
+    example['second_sentence']=example['answer0'],example['answer1'],example['answer2'],example['answer3']
+    return example
+
+test_dataset=test_dataset.map(preprocess)
+
+def tokenize(examples):
+    a=tokenizer(examples['first_sentence'],examples['second_sentence'],padding='max_length',truncation=True,max_length=256,return_tensors='jax')
+    a['labels']=examples['label']
+    return a
+
+test_dataset=test_dataset.map(tokenize)
+
+remov_col=['id', 'context', 'question', 'answer0', 'answer1', 'answer2', 'answer3', 'labels', 'context&question', 'first_sentence', 'second_sentence']
+
+test_dataset=test_dataset.remove_columns(remov_col)
+
+model = FlaxGPT2ForMultipleChoice.from_pretrained("flax-community/gpt2-Cosmos",input_shape=(1,4,1))
+
+input_id=jnp.array(test_dataset['input_ids'])
+att_mask=jnp.array(test_dataset['attention_mask'])
+
+outputs=model(input_id,att_mask)
+
+final_output=jnp.argmax(outputs,axis=-1)
+
+logger.info(f"the predction of the test dataset : {final_output[:30]}")

src/train.py

@@ -0,0 +1,241 @@
+import jax
+print(jax.local_device_count())
+import jax.numpy as jnp
+
+import flax
+import flax.linen as nn
+from flax.training.common_utils import get_metrics,onehot,shard,shard_prng_key
+from flax.training import train_state
+from flax.metrics.tensorboard import SummaryWriter
+from flax.training import checkpoints
+
+
+import logging
+import optax
+import math
+from tqdm import tqdm
+
+from pathlib import Path
+from typing import Callable
+from itertools import chain
+from flax.metrics import tensorboard
+
+from datasets import load_dataset,load_metric
+from transformers import GPT2Config,GPT2Tokenizer
+
+from model_file  import FlaxGPT2ForMultipleChoice
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+def main():
+
+
+  tokenizer=GPT2Tokenizer.from_pretrained('gpt2',pad_token='<|endoftext|>')
+
+  dataset=load_dataset('cosmos_qa')
+
+  def preprocess(example):
+    example['context&question']=example['context']+example['question']
+    example['first_sentence']=[example['context&question'],example['context&question'],example['context&question'],example['context&question']]
+    example['second_sentence']=example['answer0'],example['answer1'],example['answer2'],example['answer3']
+    return example
+
+  train_dataset=dataset['train'].map(preprocess)
+  validation_dataset=dataset['validation'].map(preprocess)
+  test_dataset=dataset['test'].map(preprocess)
+
+  #Remove after experiment
+  len_train_dataset=25262
+  len_validation_dataset=2985
+  len_test_dataset=6963
+
+  train_dataset=train_dataset.select(range(len_train_dataset))
+  test_dataset=test_dataset.select(range(len_test_dataset))
+  validation_dataset=validation_dataset.select(range(len_validation_dataset))
+
+  #remove_cols=train_dataset.column_names
+
+  def tokenize(examples):
+    a=tokenizer(examples['first_sentence'],examples['second_sentence'],padding='max_length',truncation=True,max_length=256,return_tensors='jax')
+    a['labels']=examples['label']
+    return a
+
+  train_dataset=train_dataset.map(tokenize)
+  validation_dataset=validation_dataset.map(tokenize)
+  test_dataset=test_dataset.map(tokenize)
+
+  remov_col=['id', 'context', 'question', 'answer0', 'answer1', 'answer2', 'answer3', 'labels', 'context&question', 'first_sentence', 'second_sentence']
+
+  train_dataset=train_dataset.remove_columns(remov_col)
+  validation_dataset=validation_dataset.remove_columns(remov_col)
+  test_dataset=test_dataset.remove_columns(remov_col)
+
+  per_device_batch_size=4
+  seed=0
+  num_train_epochs=3
+  learning_rate=2e-5
+  
+
+  total_batch_size = per_device_batch_size * jax.local_device_count()
+  print('The overall batch size (both for training and eval) is', total_batch_size)
+  num_train_steps = len(train_dataset) // total_batch_size * num_train_epochs
+  num_validation_steps=len(validation_dataset)//total_batch_size*num_train_epochs
+
+  learning_rate_function = optax.linear_schedule(init_value=learning_rate, end_value=0, transition_steps=num_train_steps)
+
+  class TrainState(train_state.TrainState):
+    logits_function:Callable=flax.struct.field(pytree_node=False)
+    loss_function:Callable=flax.struct.field(pytree_node=False)
+
+  def adamw(weight_decay):
+    return optax.adamw(learning_rate=learning_rate_function,b1=0.9,b2=0.99,eps=1e-6,weight_decay=weight_decay)
+
+  decay_path=lambda p:not any(x in p for x in ['bias','LayerNorm.weight'])
+
+  def traverse(function):
+    def mask(data):
+      flat=flax.traverse_util.flatten_dict(data)
+      return flax.traverse_util.unflatten_dict({k:function(k,v) for k,v in flat.items()})
+    return mask
+
+  gradient_transformation=optax.chain(
+      optax.masked(adamw(0.0),mask=traverse(lambda path,_:decay_path(path))),
+      optax.masked(adamw(0.01),mask=traverse(lambda path,_:not decay_path(path))))
+
+  def loss_function(logits,labels):
+    logits=flax.linen.log_softmax(logits)
+    xentropy=optax.softmax_cross_entropy(logits,onehot(labels,num_classes=4))
+    return jnp.mean(xentropy)
+
+  def eval_function(logits):
+    return logits.argmax(-1)
+
+  model = FlaxGPT2ForMultipleChoice.from_pretrained('gpt2-large',input_shape=(1,4,1))
+
+  state=TrainState.create(apply_fn=model.__call__,
+                          params=model.params,
+                          tx=gradient_transformation,
+                          logits_function=eval_function,
+                          loss_function=loss_function)
+                      
+  def train_step(state,batch,dropout_rng):
+      targets=batch.pop("label")
+      dropout_rng,new_dropout_rng=jax.random.split(dropout_rng)
+      def loss_function(params):
+          logits=state.apply_fn(**batch,params=params,dropout_rng=dropout_rng,train=True)[0]
+          loss=state.loss_function(logits,targets)
+          return loss
+      grad_function=jax.value_and_grad(loss_function)
+      loss,grad=grad_function(state.params)
+      grad=jax.lax.pmean(grad,"batch")
+      new_state=state.apply_gradients(grads=grad)
+        #Added.
+      logits=new_state.apply_fn(**batch,params=new_state.params,dropout_rng=dropout_rng,train=True)[0]
+      accuracy=jnp.equal(jnp.argmax(logits,axis=-1),targets)
+      metrics=jax.lax.pmean({"loss":loss,"learning_rate":learning_rate_function(state.step),'accuracy':accuracy},axis_name="batch")
+      return new_state,metrics,new_dropout_rng
+
+  parallel_train_step = jax.pmap(train_step, axis_name="batch", donate_argnums=(0,))
+
+  def eval_step(state, batch):
+      targets=batch.pop('label')
+      logits = state.apply_fn(**batch, params=state.params, train=False)
+      loss=state.loss_function(logits,targets)
+      predictions=state.logits_function(logits)
+      eval_accuracy=jnp.equal(predictions,targets)
+      #eval_acc=jnp.equal(predictions,targets)
+      metrics=jax.lax.pmean({"loss":loss,'accuracy':eval_accuracy},axis_name="batch")
+      #return state.logits_function(logits)  #(8,4)
+      return targets,predictions,metrics
+
+  parallel_eval_step = jax.pmap(eval_step, axis_name="batch")
+
+  def glue_train_data_loader(rng,dataset,batch_size):
+    steps_per_epoch=len_train_dataset//batch_size
+    perms=jax.random.permutation(rng,len(dataset))
+    perms=perms[:steps_per_epoch*batch_size]
+    perms=perms.reshape((steps_per_epoch,batch_size))
+    for perm in perms:
+      batch=dataset[perm]
+      batch={k:jnp.array(v) for k,v in batch.items()}
+      batch=shard(batch)
+      yield batch
+
+  rng=jax.random.PRNGKey(seed)
+  dropout_rngs=jax.random.split(rng,jax.local_device_count())
+
+  def glue_eval_data_loader(dataset, batch_size):
+      for i in range(len_validation_dataset // batch_size):
+          batch = dataset[i * batch_size : (i + 1) * batch_size]
+          batch = {k: jnp.array(v) for k, v in batch.items()}
+          batch = shard(batch)
+
+          yield batch
+
+  state = flax.jax_utils.replicate(state)
+  #metrics_list = list_metrics()
+
+  actual_task = "mnli"
+  metric = load_metric('glue', "mnli")
+  actual_taskmetric = load_metric('glue', actual_task)
+
+  workdir='../results_tensorboard'
+  summary_writer = tensorboard.SummaryWriter(workdir)
+  #summary_writer.hparams(dict(GPT2Config()))
+
+  logger.info(f"***** Running training *****")
+  logger.info(f"  Num examples = {len_train_dataset}")
+  logger.info(f"  Num Epochs = {num_train_epochs}")
+  logger.info(f"  Instantaneous batch size per device = {per_device_batch_size}")
+  logger.info(f"  Total train batch size  = {total_batch_size}")
+  logger.info(f"  Total optimization steps = {num_train_steps}")
+
+  for i, epoch in enumerate(tqdm(range(1, num_train_epochs+1), desc=f"Epoch ...", position=0, leave=True)):
+    rng, input_rng = jax.random.split(rng)
+    train_acc_metrics=[]
+    train_loss_metrics=[]
+    eval_acc_metrics=[]
+    eval_loss_metrics=[]
+    # train
+    with tqdm(total=len_train_dataset // total_batch_size, desc="Training...", leave=False) as progress_bar_train:
+      for idx,batch in enumerate(glue_train_data_loader(input_rng, train_dataset, total_batch_size)):
+        state, train_metric, dropout_rngs = parallel_train_step(state, batch, dropout_rngs)
+        train_acc_metrics.append(jax.device_get(train_metric['accuracy']).mean().item())
+        train_loss_metrics.append(flax.jax_utils.unreplicate(train_metric)['loss'].item())
+        if idx%5==0:
+          summary_writer.scalar('train_loss',flax.jax_utils.unreplicate(train_metric)['loss'].item(),idx)
+          summary_writer.scalar('train_accuracy', jax.device_get(train_metric['accuracy']).mean().item(),idx)
+        if idx%20==0:
+            logger.info(f"train_step_loss{idx}:   {flax.jax_utils.unreplicate(train_metric)['loss'].item()} train_step_acc{idx}:   {jax.device_get(train_metric['accuracy']).mean().item()}")
+            
+        progress_bar_train.update(1)
+
+    # evaluate
+    with tqdm(total=len_validation_dataset // total_batch_size, desc="Evaluating...", leave=False) as progress_bar_eval:
+      for idx,batch in enumerate(glue_eval_data_loader(validation_dataset, total_batch_size)):
+          labels,predictions,eval_metric=parallel_eval_step(state, batch)
+          eval_acc_metrics.append(jax.device_get(eval_metric['accuracy']).mean().item())
+          eval_loss_metrics.append(flax.jax_utils.unreplicate(eval_metric)['loss'].item())
+          progress_bar_eval.update(1)
+          if idx%5==0:
+            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()}")
+            summary_writer.scalar('eval_loss : ', flax.jax_utils.unreplicate(eval_metric)['loss'].item(),idx)
+            summary_writer.scalar('eval_accuracy : ', jax.device_get(eval_metric['accuracy']).mean().item(),idx)
+
+    if jax.process_index() == 0:
+        params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
+    
+        model.save_pretrained(
+            '../',
+            params=params,
+            push_to_hub=True,
+            commit_message=f"Saving weights of epoch {epoch} at step {idx}",)
+          
+    logger.info(f"---------------------Epoch {epoch} done-----------------")
+    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()}")
+    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()}")
+  summary_writer.flush()
+
+if __name__ == "__main__":
+    main()