rbbt-dm 1.2.6 → 1.2.9

data/python/rbbt_dm/huggingface.py

@@ -1,32 +1,41 @@
 #{{{ LOAD MODEL
 
 def import_module_class(module, class_name):
-    exec(f"from {module} import {class_name}")
+    if (not module == None):
+        exec(f"from {module} import {class_name}")
     return eval(class_name)
 
-def load_model(task, checkpoint):
-    class_name = 'AutoModelFor' + task
-    return import_module_class('transformers', class_name).from_pretrained(checkpoint)
+def load_model(task, checkpoint, **kwargs):
+    if (":" in task):
+        module, class_name = task.split(":")
+        if (task == None):
+            module, class_name = None, module
+        return import_module_class(module, class_name).from_pretrained(checkpoint, **kwargs)
+    else:
+        class_name = 'AutoModelFor' + task
+        return import_module_class('transformers', class_name).from_pretrained(checkpoint)
 
-def load_tokenizer(task, checkpoint):
+def load_tokenizer(task, checkpoint, **kwargs):
     class_name = 'AutoTokenizer'
-    return import_module_class('transformers', class_name).from_pretrained(checkpoint)
+    return import_module_class('transformers', class_name).from_pretrained(checkpoint, **kwargs)
 
 def load_model_and_tokenizer(task, checkpoint):
     model = load_model(task, checkpoint)
     tokenizer = load_tokenizer(task, checkpoint)
     return model, tokenizer
 
-def load_model_and_tokenizer_from_directory(directory):
-    import os
-    import json
-    options_file = os.path.join(directory, 'options.json')
-    f = open(options_file, "r")
-    options = json.load(f.read())
-    f.close()
-    task = options["task"]
-    checkpoint = options["checkpoint"]
-    return load_model_and_tokenizer(task, checkpoint)
+# Not used
+
+#def load_model_and_tokenizer_from_directory(directory):
+#    import os
+#    import json
+#    options_file = os.path.join(directory, 'options.json')
+#    f = open(options_file, "r")
+#    options = json.load(f.read())
+#    f.close()
+#    task = options["task"]
+#    checkpoint = options["checkpoint"]
+#    return load_model_and_tokenizer(task, checkpoint)
 
 #{{{ SIMPLE EVALUATE
 
@@ -51,21 +60,36 @@ def load_tsv(tsv_file):
     from datasets import load_dataset
     return load_dataset('csv', data_files=[tsv_file], sep="\t")
 
+def load_json(json_file):
+    from datasets import load_dataset
+    return load_dataset('json', data_files=[json_file])
+
+def tokenize_dataset(tokenizer, dataset):
+    return dataset.map(lambda subset: subset if ("input_ids" in subset.keys()) else tokenizer(subset["text"], truncation=True), batched=True)
+
 def tsv_dataset(tokenizer, tsv_file):
     dataset = load_tsv(tsv_file)
-    tokenized_dataset = dataset.map(lambda example: tokenizer(example["text"], truncation=True) , batched=True)
-    return tokenized_dataset
+    return tokenize_dataset(tokenizer, dataset)
+
+def json_dataset(tokenizer, json_file):
+    dataset = load_json(json_file)
+    return tokenize_dataset(tokenizer, dataset)
 
 def training_args(*args, **kwargs):
     from transformers import TrainingArguments
     training_args = TrainingArguments(*args, **kwargs)
     return training_args
 
-
-def train_model(model, tokenizer, training_args, tsv_file, class_weights=None):
+def train_model(model, tokenizer, training_args, dataset, class_weights=None, **kwargs):
     from transformers import Trainer
 
-    tokenized_dataset = tsv_dataset(tokenizer, tsv_file)
+    if (isinstance(dataset, str)):
+        if (dataset.endswith('.json')):
+            tokenized_dataset = json_dataset(tokenizer, dataset)
+        else:
+            tokenized_dataset = tsv_dataset(tokenizer, dataset)
+    else:
+        tokenized_dataset = tokenize_dataset(tokenizer, dataset)
 
     if (not class_weights == None):
         import torch
@@ -86,7 +110,8 @@ def train_model(model, tokenizer, training_args, tsv_file, class_weights=None):
                 model,
                 training_args,
                 train_dataset = tokenized_dataset["train"],
-                tokenizer = tokenizer
+                tokenizer = tokenizer,
+                **kwargs
                 )
     else:
 
@@ -94,7 +119,8 @@ def train_model(model, tokenizer, training_args, tsv_file, class_weights=None):
                 model,
                 training_args,
                 train_dataset = tokenized_dataset["train"],
-                tokenizer = tokenizer
+                tokenizer = tokenizer,
+                **kwargs
                 )
 
     trainer.train()
@@ -124,10 +150,16 @@ def find_tokens_in_input(dataset, token_ids):
     return position_rows
 
 
-def predict_model(model, tokenizer, training_args, tsv_file, locate_tokens = None):
+def predict_model(model, tokenizer, training_args, dataset, locate_tokens = None):
     from transformers import Trainer
 
-    tokenized_dataset = tsv_dataset(tokenizer, tsv_file)
+    if (isinstance(dataset, str)):
+        if (dataset.endswith('.json')):
+            tokenized_dataset = json_dataset(tokenizer, dataset)
+        else:
+            tokenized_dataset = tsv_dataset(tokenizer, dataset)
+    else:
+        tokenized_dataset = tokenize_dataset(tokenizer, dataset)
 
     trainer = Trainer(
             model,
@@ -143,4 +175,3 @@ def predict_model(model, tokenizer, training_args, tsv_file, locate_tokens = Non
     else:
         return result
 
-

data/python/rbbt_dm/language_model.py

@@ -0,0 +1,70 @@
+def group_texts(examples):
+    # Concatenate all texts.
+    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
+    total_length = len(concatenated_examples[list(examples.keys())[0]])
+    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+    total_length = (total_length // block_size) * block_size
+    # Split by chunks of max_len.
+    result = {
+        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+        for k, t in concatenated_examples.items()
+    }
+    result["labels"] = result["input_ids"].copy()
+    return result
+
+def whole_word_masking_data_collator(features):
+    from transformers import default_data_collator
+    for feature in features:
+        word_ids = feature.pop("word_ids")
+
+        # Create a map between words and corresponding token indices
+        mapping = collections.defaultdict(list)
+        current_word_index = -1
+        current_word = None
+        for idx, word_id in enumerate(word_ids):
+            if word_id is not None:
+                if word_id != current_word:
+                    current_word = word_id
+                    current_word_index += 1
+                mapping[current_word_index].append(idx)
+
+        # Randomly mask words
+        mask = np.random.binomial(1, wwm_probability, (len(mapping),))
+        input_ids = feature["input_ids"]
+        labels = feature["labels"]
+        new_labels = [-100] * len(labels)
+        for word_id in np.where(mask)[0]:
+            word_id = word_id.item()
+            for idx in mapping[word_id]:
+                new_labels[idx] = labels[idx]
+                input_ids[idx] = tokenizer.mask_token_id
+        feature["labels"] = new_labels
+
+    return default_data_collator(features)
+
+if __name__ == "__main__2":
+
+    from transformers import AutoModelForMaskedLM
+    from transformers import AutoTokenizer
+    import torch
+
+    model_checkpoint = "distilbert-base-uncased"
+    model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)
+    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
+
+    text = "This is a great [MASK]."
+
+    inputs = tokenizer(text, return_tensors="pt")
+    token_logits = model(**inputs).logits
+    # Find the location of [MASK] and extract its logits
+    mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]
+    mask_token_logits = token_logits[0, mask_token_index, :]
+    # Pick the [MASK] candidates with the highest logits
+    top_5_tokens = torch.topk(mask_token_logits, 5, dim=1).indices[0].tolist()
+
+    for token in top_5_tokens:
+        print(f"'>>> {text.replace(tokenizer.mask_token, tokenizer.decode([token]))}'")
+
+
+

data/python/rbbt_dm/util.py

@@ -0,0 +1,30 @@
+import random
+import torch
+import numpy
+
+def set_seed(seed):
+    """
+    Set seed in several backends
+    """
+    random.seed(seed)
+    numpy.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+def deterministic():
+    """
+    Ensure that all operations are deterministic on GPU (if used) for
+    reproducibility
+    """
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+def device():
+    return torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+
+def data_directory():
+    from pathlib import Path
+    print(Path.home())
+

data/share/spaCy/gpu/textcat_accuracy.conf

@@ -20,7 +20,8 @@ factory = "transformer"
 
 [components.transformer.model]
 @architectures = "spacy-transformers.TransformerModel.v1"
-name = "emilyalsentzer/Bio_ClinicalBERT"
+name = "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext"
+#name = "emilyalsentzer/Bio_ClinicalBERT"
 tokenizer_config = {"use_fast": true}
 
 [components.transformer.model.get_spans]

data/test/rbbt/vector/model/huggingface/test_masked_lm.rb

@@ -0,0 +1,41 @@
+require File.join(File.expand_path(File.dirname(__FILE__)),'../../../..', 'test_helper.rb')
+require 'rbbt/vector/model/huggingface/masked_lm'
+
+class TestMaskedLM < Test::Unit::TestCase
+  def test_train_new_word
+    TmpFile.with_file do |dir|
+
+      checkpoint = "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext"
+      mlm = MaskedLMModel.new checkpoint, dir, tokenizer_args: {max_length: 16, model_max_length: 16}
+
+      mod, tokenizer = mlm.init
+      if tokenizer.vocab["[GENE]"].nil?
+        tokenizer.add_tokens("[GENE]")
+        mod.resize_token_embeddings(tokenizer.__len__)
+      end
+
+      100.times do
+        mlm.add "This [GENE] is [MASK] on tumor cells.", %w(expressed)
+        mlm.add "This [MASK] is expressed.", %w([GENE])
+      end
+
+      assert_equal "protein", mlm.eval(["This [MASK] is expressed."])
+
+      mlm.train
+
+      assert_equal "[GENE]", mlm.eval(["This [MASK] is expressed."])
+      assert_equal "expressed", mlm.eval(["This [GENE] is [MASK] in tumor cells."])
+
+      mlm = MaskedLMModel.new checkpoint, dir, :max_length => 16
+      
+      assert_equal "[GENE]", mlm.eval(["This [MASK] is expressed."])
+      assert_equal "expressed", mlm.eval(["This [GENE] is [MASK] in tumor cells."])
+
+      mlm = VectorModel.new dir
+      
+      assert_equal "[GENE]", mlm.eval(["This [MASK] is expressed."])
+      assert_equal "expressed", mlm.eval(["This [GENE] is [MASK] in tumor cells."])
+
+    end
+  end
+end