:bugs: fix some bugs

Author: AntoineSimoulin

.gitignore

@@ -0,0 +1 @@
+model/

README.md

@@ -5,7 +5,7 @@
 <p>
 
 **PyTree** implements tree-structured neural networks in PyTorch.
-The package provides highly generic recursive neural network implementations as well as efficient batching methods.
+The package provides highly generic tree-structured neural network implementations as well as efficient batching methods.
 
 ## Installation
 

examples/README.md

@@ -14,6 +14,10 @@ python pytree/examples/run_sick.py \
     --num_train_epochs 20
 ```
 
+```
+CUDA_VISIBLE_DEVICES=2 python examples/run_sick.py     --glove_file_path /data/asimouli/GLOVE/glove.6B.300d.txt     --do_train     --do_eval      --output_dir './model'     --dataset_name 'sick'     --remove_unused_columns False     --learning_rate 0.05     --per_device_train_batch_size 25     --num_train_epochs 15    --weight_decay 1e-4  --lr_scheduler_type constant  --do_predict    --overwrite_cache True  --overwrite_output_dir
+```
+
 ## References
 
 > <div id="tai-2015">Kai Sheng Tai, Richard Socher, Christopher D. Manning <a href=https://aclanthology.org/P15-1150>Improved Semantic Representations From Tree-Structured Long Short-Term Memory Networks.</a> ACL (1) 2015: 1556-1566</div>

examples/main.py

@@ -0,0 +1,227 @@
+import logging
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = ""
+import sys
+from dataclasses import dataclass, field
+from typing import Optional
+
+import datasets
+from datasets import load_dataset, load_metric
+
+import transformers
+from transformers import (
+    Trainer,
+    EvalPrediction,
+    HfArgumentParser,
+    TrainingArguments,
+    set_seed,
+    default_data_collator,
+)
+from transformers.trainer_utils import get_last_checkpoint
+from transformers.utils import check_min_version
+from transformers.utils.versions import require_version
+# from utils_qa import postprocess_qa_predictions
+
+from pytree import (
+    NaryConfig, 
+    NaryTree,
+    ChildSumConfig,
+    ChildSumTree,
+    GloveTokenizer,
+    Similarity,
+    SimilarityConfig
+)
+from pytree.data import prepare_input_from_constituency_tree, prepare_input_from_dependency_tree
+from pytree.data.utils import build_tree_ids_n_ary
+
+from supar import Parser
+import torch
+import numpy as np
+import math
+from sklearn.metrics import mean_squared_error
+from scipy.stats import pearsonr, spearmanr
+
+# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
+check_min_version("4.11.0")  # 4.12.0.dev0
+
+require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt")
+
+logger = logging.getLogger(__name__)
+
+class SickTrainer(Trainer):
+    
+    def create_optimizer(self):
+        """
+        Setup the optimizer.
+
+        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
+        Trainer's init through :obj:`optimizers`, or subclass and override this method in a subclass.
+        """
+        self.optimizer = torch.optim.Adagrad(self.model.parameters(), lr=0.025, weight_decay=self.args.weight_decay)
+    
+        # if self.sharded_ddp == ShardedDDPOption.SIMPLE:
+        #     self.optimizer = OSS(
+        #         params=optimizer_grouped_parameters,
+        #         optim=optimizer_cls,
+        #         **optimizer_kwargs,
+        #     )
+        # else:
+        #     self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
+
+        # if is_sagemaker_mp_enabled():
+        #     self.optimizer = smp.DistributedOptimizer(self.optimizer)
+
+        return self.optimizer
+
+
+con = Parser.load('crf-con-en')
+glove_tokenizer = GloveTokenizer(glove_file_path='/data/asimouli/GLOVE/glove.6B.300d.txt', vocab_size=10000)
+
+config = NaryConfig()
+encoder = NaryTree(config)
+encoder.embeddings.load_pretrained_embeddings(
+    torch.tensor(glove_tokenizer.embeddings_arr))    
+config_similarity = SimilarityConfig()
+model = Similarity(encoder, config_similarity)
+
+raw_datasets = load_dataset('sick')
+column_names = raw_datasets["train"].column_names
+
+def map_label_to_target(label, num_classes):
+    target = [0] * num_classes  # torch.zeros(1, num_classes, dtype=torch.float)
+    ceil = int(math.ceil(label))
+    floor = int(math.floor(label))
+    if ceil == floor:
+        target[floor - 1] = 1
+    else:
+        target[floor - 1] = ceil - label
+        target[ceil - 1] = label - floor
+    return target
+
+def prepare_train_features(examples):
+    examples['input_ids_A'] = []
+    examples['input_ids_B'] = []
+    examples['head_idx_A'] = []
+    examples['head_idx_B'] = []
+    examples['labels'] = []
+    
+    for sent_A in examples['sentence_A']:
+        con_tree_A = str(con.predict(sent_A.split(), verbose=False)[0])
+        input_ids_A, head_idx_A = prepare_input_from_constituency_tree(con_tree_A)
+        input_ids_A = glove_tokenizer.convert_tokens_to_ids(input_ids_A)
+        examples['input_ids_A'].append(input_ids_A)
+        examples['head_idx_A'].append(head_idx_A)
+    
+    for sent_B in examples['sentence_B']:
+        con_tree_B = str(con.predict(sent_B.split(), verbose=False)[0])
+        input_ids_B, head_idx_B = prepare_input_from_constituency_tree(con_tree_B)
+        input_ids_B = glove_tokenizer.convert_tokens_to_ids(input_ids_B)
+        examples['input_ids_B'].append(input_ids_B)
+        examples['head_idx_B'].append(head_idx_B)
+
+    for rel_score in examples['relatedness_score']:
+        examples['labels'].append(map_label_to_target(rel_score, 5))
+
+    return examples
+
+training_args = TrainingArguments(
+    learning_rate=0.025, 
+    per_device_train_batch_size=25, 
+    num_train_epochs=20, 
+    weight_decay=1e-4, 
+    lr_scheduler_type='constant', 
+    output_dir="/home/asimouli/PhD/PyTree/pytree_remote/model", 
+    do_train=True, 
+    do_eval=True,
+    remove_unused_columns=False)
+
+train_examples = raw_datasets["train"]
+with training_args.main_process_first(desc="train dataset map pre-processing"):
+    train_dataset = train_examples.map(
+        prepare_train_features,
+        batched=True,
+        num_proc=None,
+        remove_columns=None,
+        load_from_cache_file=True,
+        desc="Running parser on train dataset",
+    )
+
+# # Validation preprocessing
+
+eval_examples = raw_datasets["validation"]
+eval_dataset = eval_examples.map(
+    prepare_train_features,
+    batched=True,
+    num_proc=None,
+    remove_columns=None,  # column_names,
+    desc="Running parser on validation dataset",
+)
+
+def data_collator_with_padding(features, pad_ids=0, columns=None):
+    batch = {}
+    first = features[0]
+    if columns is None:
+        columns = ["head_idx_A", "head_idx_B", "input_ids_A", "input_ids_B"]
+    feature_max_len = {k: max([len(f[k]) for f in features]) for k in first.keys() if k in columns or len(columns) == 0}
+    for k, v in first.items():
+        if k in columns or len(columns) == 0:
+            feature_padded = [list([int(ff) for ff in f[k]]) + [0] * (feature_max_len[k] - len(f[k])) for f in features]
+            batch[k] = feature_padded  # [f[k] for f in features]
+    tree_ids_A, tree_ids_r_A, tree_ids_l_A = build_tree_ids_n_ary(batch['head_idx_A'])
+    tree_ids_B, tree_ids_r_B, tree_ids_l_B = build_tree_ids_n_ary(batch['head_idx_B'])
+    batch['input_ids_A'] = torch.tensor(batch['input_ids_A'])
+    batch['input_ids_B'] = torch.tensor(batch['input_ids_B'])
+    batch['tree_ids_A'] = torch.tensor(tree_ids_A)
+    batch['tree_ids_B'] = torch.tensor(tree_ids_B)
+    batch['tree_ids_r_A'] = torch.tensor(tree_ids_r_A)
+    batch['tree_ids_r_B'] = torch.tensor(tree_ids_r_B)
+    batch['tree_ids_l_A'] = torch.tensor(tree_ids_l_A)
+    batch['tree_ids_l_B'] = torch.tensor(tree_ids_l_B)
+    batch['labels'] = torch.tensor([f['labels'] for f in features])
+    return batch
+
+data_collator = data_collator_with_padding
+
+def compute_metrics(eval_prediction):
+    prediction = np.matmul(np.exp(eval_prediction.predictions), np.arange(1, 5 + 1))
+    target = np.matmul(eval_prediction.label_ids, np.arange(1, 5 + 1))
+    results_relatedness = {
+        'pearson': pearsonr(prediction, target)[0] * 100,
+        'spearman': spearmanr(prediction, target)[0] * 100,
+        'mse': mean_squared_error(prediction, target) * 100
+    }
+    return results_relatedness
+    
+trainer = SickTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    data_collator=data_collator,
+    compute_metrics=compute_metrics,
+    optimizers=("Adagrad", None),
+)
+
+# Training
+
+train_result = trainer.train(resume_from_checkpoint=None)
+trainer.save_model()  # Saves the tokenizer too for easy upload
+
+metrics = train_result.metrics
+max_train_samples = len(train_dataset)
+metrics["train_samples"] = min(max_train_samples, len(train_dataset))
+
+trainer.log_metrics("train", metrics)
+trainer.save_metrics("train", metrics)
+trainer.save_state()
+
+
+logger.info("*** Evaluate ***")
+metrics = trainer.evaluate()
+
+max_eval_samples = len(eval_dataset)
+metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))
+
+trainer.log_metrics("eval", metrics)
+trainer.save_metrics("eval", metrics)
+

examples/run_sick.py

@@ -67,6 +67,33 @@
 
 logger = logging.getLogger(__name__)
 
+class SickTrainer(Trainer):
+    
+    def create_optimizer(self):
+        """
+        Setup the optimizer.
+
+        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
+        Trainer's init through :obj:`optimizers`, or subclass and override this method in a subclass.
+        """
+        self.optimizer = torch.optim.Adagrad(self.model.parameters(), 
+            lr=self.args.learning_rate, 
+            weight_decay=self.args.weight_decay)
+    
+        # if self.sharded_ddp == ShardedDDPOption.SIMPLE:
+        #     self.optimizer = OSS(
+        #         params=optimizer_grouped_parameters,
+        #         optim=optimizer_cls,
+        #         **optimizer_kwargs,
+        #     )
+        # else:
+        #     self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
+
+        # if is_sagemaker_mp_enabled():
+        #     self.optimizer = smp.DistributedOptimizer(self.optimizer)
+
+        return self.optimizer
+
 
 @dataclass
 class ModelArguments:
@@ -278,7 +305,7 @@ def main():
     if data_args.dataset_name is not None:
         # Downloading and loading a dataset from the hub.
         raw_datasets = load_dataset(
-            data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir
+            data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir, keep_in_memory=False,
         )
     else:
         data_files = {}
@@ -303,7 +330,7 @@ def main():
     # download model & vocab.
     # dep = Parser.load('biaffine-dep-en')
     con = Parser.load('crf-con-en')
-    glove_tokenizer = GloveTokenizer(glove_file_path=data_args.glove_file_path, vocab_size=10000)
+    glove_tokenizer = GloveTokenizer(glove_file_path=data_args.glove_file_path, vocab_size=100000)
     # config = ChildSumConfig()
     # encoder = ChildSumTree(config)
     config = NaryConfig()
@@ -577,15 +604,15 @@ def compute_metrics(eval_prediction):
     #     return metric.compute(predictions=p.predictions, references=p.label_ids)
 
     # Initialize our Trainer
-    trainer = Trainer(
+    trainer = SickTrainer(
         model=model,
         args=training_args,
         train_dataset=train_dataset if training_args.do_train else None,
-        eval_dataset=eval_dataset if training_args.do_eval else None,
+        eval_dataset=predict_dataset if training_args.do_eval else None,  # eval_dataset
         # eval_examples=eval_examples if training_args.do_eval else None,
         data_collator=data_collator,
         compute_metrics=compute_metrics,
-        # optimizers=(torch.optim.Adagrad(model.parameters(), weight_decay=1e-4), None),
+        optimizers=("Adagrad", None),
     )
     # trainer = QuestionAnsweringTrainer(
     #     model=model,

pytree/data/utils.py

@@ -82,11 +82,14 @@ def get_nodes_detph(head_idx):
   return depth
 
 
+# def pad_tree_ids(tree_ids, depth):
+#   tree_depth = tree_ids.shape[0]
+#   padding = np.zeros((max(depth - tree_depth, 0), tree_ids.shape[1]), dtype=tree_ids.dtype)
+#   return np.concatenate((tree_ids, padding), axis=0)
 def pad_tree_ids(tree_ids, depth):
   tree_depth = tree_ids.shape[0]
   padding = np.zeros((max(depth - tree_depth, 0), tree_ids.shape[1]), dtype=tree_ids.dtype)
-  return np.concatenate((tree_ids, padding), axis=0)
-
+  return np.concatenate((padding, tree_ids), axis=0)
 
 # def build_tree_ids(head_idx):
 #   if isinstance(head_idx[0], list):
@@ -125,12 +128,18 @@ def build_tree_ids_n_ary(head_idx):
       np.array([pad_tree_ids(t[2], depth) for t in tree_ids])
   tree_ids = []
   node_idx = [get_root(head_idx)]
+  # while len(node_idx) > 0:
+  #   node_idx = get_childrens(node_idx, head_idx)
+  #   tree_step = [h_idx if idx in node_idx else 0 for idx, h_idx in enumerate(head_idx)]
+  #   tree_ids.append(tree_step)
+  # tree_ids = tree_ids[:-1]
+  # tree_ids.append(range(0, len(head_idx)))
   while len(node_idx) > 0:
     node_idx = get_childrens(node_idx, head_idx)
     tree_step = [h_idx if idx in node_idx else 0 for idx, h_idx in enumerate(head_idx)]
-    tree_ids.append(tree_step)
-  tree_ids = tree_ids[:-1]
-  tree_ids.append(range(0, len(head_idx)))
+    tree_ids.insert(0, tree_step)
+  tree_ids = tree_ids[1:]
+  tree_ids.insert(0, range(0, len(head_idx)))
   tree_ids_r = [[t if (i % 2 == 0) else 0 for (i, t) in enumerate(ti)] for ti in tree_ids]
   tree_ids_d = [[t if (i % 2 == 1) else 0 for (i, t) in enumerate(ti)] for ti in tree_ids]
   return np.array(tree_ids), np.array(tree_ids_r), np.array(tree_ids_d)