Update README.md

README.md

@@ -26,20 +26,20 @@ library_name: transformers
 
 ## Model description
 
-This model was created during the research collaboration between Bayer Pharma and Serbian Institute for Artificial Intelligence Research and Development. 
-The model is trained on about 25+ biomedical NER classes and can perform also zero-shot inference and can be further fine-tuned for new classes with just few examples (few-shot learning). 
-For more details about our methods please see the paper named ["A transformer-based method for zero and few-shot biomedical named entity recognition"](https://arxiv.org/abs/2305.04928). The model corresponds to BioBERT-based mode, trained with 1 in the first segment (check paper for more details).
+This model was created during the research collaboration between Bayer Pharma and The Institute for Artificial Intelligence Research and Development of Serbia. 
+The model is trained on 26 biomedical Named Entity (NE) classes and can perform zero-shot inference. It also can be further fine-tuned for new classes with just few examples (few-shot learning). 
+For more details about our method please see the paper named ["A transformer-based method for zero and few-shot biomedical named entity recognition"](https://arxiv.org/abs/2305.04928). The model corresponds to PubMedBERT-based model, trained with 1 in the first segment (check paper for more details).
 
-Model takes as input two strings. String1 is NER label that is being searched in second string. String1 must be phrase for entity. String2 is short text where String1 is searched for semantically.
-model outputs list of zeros and ones corresponding to the occurance of Named Entity and corresponing to the tokens(tokens given by transformer tokenizer) of the Sring2.
+Model takes two strings as input. String1 is NE label that is being searched in second string. String2 is short text where one wants to searc for NE (represented by String1).
+Model outputs list of ones (corresponding to the found Named Entities) and zeros (corresponding to other non-NE tokens) of the Sring2.
 
 ## Example of usage
 ```python
 from transformers import AutoTokenizer
 from transformers import BertForTokenClassification
 
-modelname = 'ProdicusII/ZeroShotBioNER'  # modelpath
-tokenizer = AutoTokenizer.from_pretrained(modelname)  ## loading the tokenizer of that model
+modelname = 'MilosKorsRad/BioNER'  # modelpath
+tokenizer = AutoTokenizer.from_pretrained(modelname)  ## loading the tokenizer of the model
 string1 = 'Drug'
 string2 = 'No recent antibiotics or other nephrotoxins, and no symptoms of UTI with benign UA.'
 encodings = tokenizer(string1, string2, is_split_into_words=False,
@@ -53,34 +53,42 @@ print(prediction_logits)
 
 ## Example of fine-tuning with few-shot learning
 
-In order to fine-tune model to the new entity using few shots, the dataset needs to be transformed to torch.utils.data.Dataset, containing BERT tokens and set of 0s and 1s (1 is where the class is positive and should be predicted as the member of given NER class). After the dataset is created, the following can be done (for more details, please have a look at the code at GitHub - https://github.com/br-ai-ns-institute/Zero-ShotNER): 
+In order to fine-tune model with new entity using few-shots, the dataset needs to be transformed to torch.utils.data.Dataset, containing BERT tokens and set of 0s and 1s (1 is where the class is positive and should be predicted as the member of given NE class). After the dataset is created, the following can be done (for more details, please have a look at the code at GitHub - https://github.com/br-ai-ns-institute/Zero-ShotNER): 
 
 ```python
- training_args = TrainingArguments(
-        output_dir=os.path.join('Results', class_unseen, str(j)+'Shot'),  # folder for results
-        num_train_epochs=10,                                              # number of epochs
-        per_device_train_batch_size=16,                                   # batch size per device during training
-        per_device_eval_batch_size=16,                                    # batch size for evaluation
-        weight_decay=0.01,                                                # strength of weight decay
-        logging_dir=os.path.join('Logs', class_unseen, str(j)+'Shot'),    # folder for logs
+ for i in [train1shot, train10shot, train100shot]:
+    training_args = TrainingArguments(
+        output_dir='./Results'+class_unseen+'FewShot'+str(i),  # output folder (folder to store the results)
+        num_train_epochs=10,                                   # number of training epochs
+        per_device_train_batch_size=16,                        # batch size per device during training
+        per_device_eval_batch_size=16,                         # batch size for evaluation
+        weight_decay=0.01,                                     # strength of weight decay
+        logging_dir='./Logs'+class_unseen+'FewShot'+str(i),    # folder to store the logs
         save_strategy='epoch',
         evaluation_strategy='epoch',
-        load_best_model_at_end=True,
+        load_best_model_at_end=True 
     )
-    
-model0 = BertForTokenClassification.from_pretrained(model_path, num_labels=2)
-trainer = Trainer(
-    model=model0,                # pretrained model
-    args=training_args,          # training artguments
-    train_dataset=dataset,       # Object of class torch.utils.data.Dataset for training
-    eval_dataset=dataset_valid   # Object of class torch.utils.data.Dataset for vaLidation
+
+    model0 = BertForTokenClassification.from_pretrained(model_path, num_labels=2)
+
+    trainer = Trainer(
+        model=model0,                # pre-trained model for fine-tuning
+        args=training_args,          # training arguments defined above
+        train_dataset=train_0shot,   # dataset class object for training
+        eval_dataset=valid_dataset   # dataset class object for validation
     )
-start_time = time.time()
-trainer.train()
-total_time = time.time()-start_time
-model0_path = os.path.join('Results', class_unseen, str(j)+'Shot', 'Model')
-os.makedirs(model0_path, exist_ok=True)
-trainer.save_model(model0_path)
+
+    start_time = time.time()
+    trainer.train()
+    total_time = time.time()-start_time
+
+    model_path = os.path.join('Results', class_unseen, 'FewShot',str(i), 'Model')
+    os.makedirs(model_path, exist_ok=True)
+    model.save_pretrained(model_path)
+
+    tokenizer_path = os.path.join('Results', class_unseen, 'FewShot', str(i), 'Tokenizer')
+    os.makedirs(tokenizer_path, exist_ok=True)
+    tokenizer.save_pretrained(tokenizer_path)
 ```
 
 ## Available classes
@@ -123,7 +131,7 @@ The following datasets and entities were used for training and therefore they ca
   * ADE 
   * Duration
 
-On top of this, one can use the model in zero-shot regime with other classes, and also fine-tune it with few examples of other classes. 
+On top of this, one can use the model for zero-shot learning with other classes, and also fine-tune it with few examples of other classes.