RelCAT: flake8 fixes.

Author: vladd-bit

medcat/utils/relation_extraction/ml_utils.py

@@ -23,6 +23,7 @@ def split_list_train_test_by_class(data: List, sample_limit: int = -1, test_size
     Args:
         data (List): "output_relations": relation_instances, <-- see create_base_relations_from_doc/csv
                     for data columns
+        sample_limit (int, optional): limit the number of samples per class, useful for dataset balancing . Defaults to -1.
         test_size (float): Defaults to 0.2.
         shuffle (bool): shuffle data randomly. Defaults to True.
 

medcat/utils/relation_extraction/models.py

@@ -174,7 +174,7 @@ def __init__(self, pretrained_model_name_or_path: str, relcat_config: ConfigRelC
                     this can be a HF model i.e: "bert-base-uncased", if left empty, it is normally assumed that a model is loaded from 'model.dat'
                     using the RelCAT.load() method. So if you are initializing/training a model from scratch be sure to base it on some model.
             relcat_config (ConfigRelCAT): relcat config.
-            model_config (BertConfig): HF bert config for model.
+            model_config (ModernBertConfig): HF bert config for model.
         """
         super(ModernBertModel_RelationExtraction, self).__init__()
 
@@ -268,7 +268,6 @@ def output2logits(self, sequence_output: torch.Tensor, input_ids: torch.Tensor,
         """
 
         Args:
-            pooled_output (torch.Tensor): embedding of the CLS token
             sequence_output (torch.Tensor): hidden states/embeddings for each token in the input text
             input_ids (torch.Tensor): input token ids.
             e1_e2_start (torch.Tensor): annotation tags token position

medcat/utils/relation_extraction/rel_dataset.py

@@ -192,8 +192,8 @@ def _create_relation_validation(self,
             ent2_start_char_pos (int): ent2 start char pos
             ent1_end_char_pos (int): ent1 end char pos
             ent2_end_char_pos (int): ent2 end char pos
-            ent1_token_start_pos (int, optional): ent1_token_start_pos. Defaults to -1.
-            ent2_token_start_pos (int, optional): ent2_token_start_pos. Defaults to -1.
+            ent1_token_start_pos (int): ent1_token_start_pos. Defaults to -1.
+            ent2_token_start_pos (int): ent2_token_start_pos. Defaults to -1.
             ent1_token_end_pos (int, optional): ent1_token_end_pos. Defaults to -1.
             ent2_token_end_pos (int, optional): ent2_token_end_pos. Defaults to -1.
             is_spacy_doc (bool, optional): checks if doc is spacy docs. Defaults to False.