Update README.md

README.md

@@ -36,14 +36,16 @@ dress him every morning?"],
 ```bash
 # Installing requirements
 pip install transformers
+pip install datasets 
 ```
 
 ```python
 from model_file import FlaxGPT2ForMultipleChoice
+from datasets import Dataset
 model_path="flax-community/gpt2-Cosmos"
 model = FlaxGPT2ForMultipleChoice.from_pretrained(model_path,input_shape=(1,4,1))
 
-run_dataset=Dataset.from_csv('......')
+dataset=Dataset.from_csv('......')
 
 def preprocess(example):
     example['context&question']=example['context']+example['question']
@@ -51,17 +53,17 @@ def preprocess(example):
     example['second_sentence']=example['answer0'],example['answer1'],example['answer2'],example['answer3']
     return example
     
-run_dataset=run_dataset.map(preprocess)
+dataset=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)
+dataset=dataset.map(tokenize)
 
-input_id=jnp.array(run_dataset['input_ids'])
-att_mask=jnp.array(run_dataset['attention_mask'])
+input_id=jnp.array(dataset['input_ids'])
+att_mask=jnp.array(dataset['attention_mask'])
 
 outputs=model(input_id,att_mask)