list

src/test_hellaswag.py

@@ -39,6 +39,7 @@ def tokenize(examples):
 test_dataset=test_dataset.map(tokenize)
 
 test_dataset=test_dataset.remove_columns(remove_col)
+list1=[]
 
 def glue_test_data_loader(rng,dataset,batch_size):
   steps_per_epoch=len_test_dataset//batch_size
@@ -46,10 +47,11 @@ def glue_test_data_loader(rng,dataset,batch_size):
   perms=perms[:steps_per_epoch*batch_size]
   perms=perms.reshape((steps_per_epoch,batch_size))
   for perm in perms:
+    list1.append(perm)
     batch=dataset[perm]
     #print(jnp.array(batch['label']))
     batch={k:jnp.array(v) for k,v in batch.items()}
-    batch=shard(batch)
+    #batch=shard(batch)
     yield batch
 
 seed=0
@@ -59,7 +61,7 @@ dropout_rngs=jax.random.split(rng,jax.local_device_count())
 input_id=jnp.array(test_dataset['input_ids'])
 att_mask=jnp.array(test_dataset['attention_mask'])
 
-total_batch_size=32
+total_batch_size=16
 
 from  model_file  import FlaxGPTNeoForMultipleChoice
 
@@ -69,12 +71,10 @@ restored_output=[]
 rng, input_rng = jax.random.split(rng)
 for idx,batch in enumerate(glue_test_data_loader(input_rng, test_dataset, total_batch_size)):
     outputs=model(batch['input_ids'],batch['attention_mask'])
-    #outputs=outputs['logits'].reshape(total_batch_size,-1)
-    print(outputs.shape)
     final_output=jnp.argmax(outputs,axis=-1)
     restored_output.append(final_output)
 
-finall=pd.DataFrame({'predictions':restored_output})
+finall=pd.DataFrame({'predictions':restored_output,'permutation':list1})
 finall.to_csv('../predictions.csv')