lints the code

.pylintrc

@@ -428,7 +428,8 @@ disable=raw-checker-failed,
         suppressed-message,
         useless-suppression,
         deprecated-pragma,
-        use-symbolic-message-instead
+        use-symbolic-message-instead,
+	R0902
 
 # Enable the message, report, category or checker with the given id(s). You can
 # either give multiple identifier separated by comma (,) or put this option

elise/src/configs/__init__.py

@@ -1 +1,4 @@
+"""
+All configs for ML project
+"""
 from .train_t5 import T5TrainingConfig

elise/src/configs/logging_config.yaml

@@ -1,12 +1,14 @@
 version: 1
+disable_existing_loggers: False
 formatters:
   simple:
     format: '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 handlers:
   console:
+    level: DEBUG
     class: logging.StreamHandler
     formatter: simple
     stream: ext://sys.stdout
-Root:
+root:
   Level: DEBUG
   handlers: [console]

elise/src/configs/train_t5.py

@@ -1,3 +1,7 @@
+"""
+Training config for T5 Seq2Seq training
+"""
+
 from dataclasses import dataclass
 
 

elise/src/data/mit_seq2seq_dataset.py

@@ -87,6 +87,9 @@ class MITRestaurants:
 
 
 def get_default_transforms():
+    """
+    Default transformation to convert ner dataset to seq2seq
+    """
     label_names = {v: k for k, v in MITRestaurants.ner_tags.items()}
 
     def decode_tags(tags, words):

elise/src/excutors/trainer_seq2seq.py

@@ -1,208 +0,0 @@
-from transformers import get_scheduler
-import torch
-from torch.utils.data import DataLoader
-from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModelForTokenClassification
-from transformers import DataCollatorForTokenClassification
-from accelerate import Accelerator
-import evaluate
-import datasets
-
-from tqdm.auto import tqdm
-
-
-ner_tags = {
-    "O": 0,
-    "B-Rating": 1,
-    "I-Rating": 2,
-    "B-Amenity": 3,
-    "I-Amenity": 4,
-    "B-Location": 5,
-    "I-Location": 6,
-    "B-Restaurant_Name": 7,
-    "I-Restaurant_Name": 8,
-    "B-Price": 9,
-    "B-Hours": 10,
-    "I-Hours": 11,
-    "B-Dish": 12,
-    "I-Dish": 13,
-    "B-Cuisine": 14,
-    "I-Price": 15,
-    "I-Cuisine": 16,
-}
-
-
-label_names = {v: k for k, v in ner_tags.items()}
-
-# dataset aggregation 
-dataset = load_dataset("tner/mit_restaurant")
-dataset["train"] = datasets.concatenate_datasets([dataset["train"], dataset["validation"]])
-dataset["train"] = datasets.concatenate_datasets([dataset["train"], dataset["test"]])
-
-print(dataset)
-
-
-tokenizer = AutoTokenizer.from_pretrained(
-    'sentence-transformers/all-MiniLM-L6-v2')
-
-
-def align_labels_with_tokens(labels, word_ids):
-    new_labels = []
-    current_word = None
-    for word_id in word_ids:
-        if word_id != current_word:
-            # Start of a new word!
-            current_word = word_id
-            label = -100 if word_id is None else labels[word_id]
-            new_labels.append(label)
-        elif word_id is None:
-            # Special token
-            new_labels.append(-100)
-        else:
-            # Same word as previous token
-            label = labels[word_id]
-            # If the label is B-XXX we change it to I-XXX
-            label_name = label_names[label]
-            if label_name.startswith("B"):
-                label = ner_tags["I" + label_name[1:]]
-            new_labels.append(label)
-
-    return new_labels
-
-
-def tokenize_and_align_labels(examples):
-    tokenized_inputs = tokenizer(
-        examples["tokens"], truncation=True, is_split_into_words=True
-    )
-    all_labels = examples["tags"]
-    new_labels = []
-    for i, labels in enumerate(all_labels):
-        word_ids = tokenized_inputs.word_ids(i)
-        new_labels.append(align_labels_with_tokens(labels, word_ids))
-
-    tokenized_inputs["labels"] = new_labels
-    return tokenized_inputs
-
-
-tokenized_datasets = dataset.map(
-    tokenize_and_align_labels,
-    batched=True,
-    remove_columns=dataset["train"].column_names,
-)
-
-
-def train():
-    metric = evaluate.load("seqeval")
-    data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)
-
-    train_dataloader = DataLoader(
-        tokenized_datasets["train"],
-        shuffle=True,
-        collate_fn=data_collator,
-        batch_size=128,
-    )
-    eval_dataloader = DataLoader(
-        tokenized_datasets["test"],
-        collate_fn=data_collator,
-        batch_size=8
-    )
-
-    model = AutoModelForTokenClassification.from_pretrained(
-        'sentence-transformers/all-MiniLM-L6-v2',
-        id2label=label_names,
-        label2id=ner_tags,
-    )
-
-    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
-
-    accelerator = Accelerator()
-    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
-        model, optimizer, train_dataloader, eval_dataloader
-    )
-
-    num_train_epochs = 50
-    num_update_steps_per_epoch = len(train_dataloader)
-    num_training_steps = num_train_epochs * num_update_steps_per_epoch
-
-    lr_scheduler = get_scheduler(
-        "linear",
-        optimizer=optimizer,
-        num_warmup_steps=0,
-        num_training_steps=num_training_steps,
-    )
-
-    def postprocess(predictions, labels):
-        predictions = predictions.detach().cpu().clone().numpy()
-        labels = labels.detach().cpu().clone().numpy()
-
-        # Remove ignored index (special tokens) and convert to labels
-        true_labels = [[label_names[l] for l in label if l != -100]
-                       for label in labels]
-        true_predictions = [
-            [label_names[p] for (p, l) in zip(prediction, label) if l != -100]
-            for prediction, label in zip(predictions, labels)
-        ]
-        return true_labels, true_predictions
-
-    progress_bar = tqdm(range(num_training_steps))
-
-    for epoch in range(num_train_epochs):
-        # Training
-        model.train()
-        for batch in train_dataloader:
-            outputs = model(**batch)
-            loss = outputs.loss
-            accelerator.backward(loss)
-
-            optimizer.step()
-            lr_scheduler.step()
-            optimizer.zero_grad()
-            progress_bar.update(1)
-
-        # Evaluation
-        model.eval()
-        for batch in eval_dataloader:
-            with torch.no_grad():
-                outputs = model(**batch)
-
-            predictions = outputs.logits.argmax(dim=-1)
-            labels = batch["labels"]
-
-            # Necessary to pad predictions and labels for being gathered
-            predictions = accelerator.pad_across_processes(
-                predictions, dim=1, pad_index=-100)
-            labels = accelerator.pad_across_processes(
-                labels, dim=1, pad_index=-100)
-
-            predictions_gathered = accelerator.gather(predictions)
-            labels_gathered = accelerator.gather(labels)
-
-            true_predictions, true_labels = postprocess(
-                predictions_gathered, labels_gathered)
-            metric.add_batch(predictions=true_predictions,
-                             references=true_labels)
-
-        results = metric.compute()
-        print(
-            f"epoch {epoch}:",
-            {
-                key: results[f"overall_{key}"]
-                for key in ["precision", "recall", "f1", "accuracy"]
-            },
-        )
-
-        output_dir = "restaurant_ner"
-        # Save and upload
-        accelerator.wait_for_everyone()
-        unwrapped_model = accelerator.unwrap_model(model)
-        unwrapped_model.save_pretrained(
-            output_dir, save_function=accelerator.save)
-        if accelerator.is_main_process:
-            tokenizer.save_pretrained(output_dir)
-
-    accelerator.wait_for_everyone()
-    unwrapped_model = accelerator.unwrap_model(model)
-    unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
-
-
-train()

elise/src/train_t5_seq2seq.py

@@ -1,22 +1,27 @@
+"""
+Training Flant_T5 model on tner/mit_restaurant on seq2seq task
+"""
+from dataclasses import asdict
+
 import torch
 import evaluate
 import datasets
 from torch.utils.data import DataLoader
-from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from dataclasses import asdict
-
-from transformers import DataCollatorForSeq2Seq
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
+    DataCollatorForSeq2Seq,
+    get_scheduler,
+)
 from accelerate import Accelerator
-from transformers import get_scheduler
 import numpy as np
 import mlflow
 
 from tqdm.auto import tqdm
 
-from data import MITRestaurants, get_default_transforms
 from utils.logger import get_logger
 from configs import T5TrainingConfig
+from data import MITRestaurants, get_default_transforms
 
 log = get_logger("Flan_T5")
 log.debug("heloooooooooooo?")
@@ -36,6 +41,7 @@ model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
 
 
 def tokenize(example):
+    """Tokenizes dataset for seq2seq task"""
     tokenized = tokenizer(
         example["tokens"],
         text_target=example["labels"],
@@ -57,6 +63,7 @@ metric = evaluate.load("sacrebleu")
 
 
 def postprocess(predictions, labels):
+    """Post processing to convert model output for evaluation"""
     predictions = predictions.cpu().numpy()
     labels = labels.cpu().numpy()
 
@@ -115,7 +122,8 @@ model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
 progress_bar = tqdm(range(num_training_steps))
 
 
-def train(model, dataset, metric):
+def train():
+    """Training function for finetuing flanT5"""
     # log.info("Starting Training")
     print("Starting Traning")
     for epoch in range(config.epochs):
@@ -174,10 +182,11 @@ def train(model, dataset, metric):
         transformers_model={"model": unwrapped_model, "tokenizer": tokenizer},
         task="text2text-generation",
         artifact_path="seq2seq_model",
-        registered_model_name="FlanT5_MIT"
+        registered_model_name="FlanT5_MIT",
     )
 
+
 mlflow.set_tracking_uri("http://127.0.0.1:5000")
 with mlflow.start_run() as mlflow_run:
     mlflow.log_params(asdict(config))
-    train(model, tokenized_datasets, metric)
+    train()

elise/src/utils/logger.py

@@ -4,7 +4,9 @@ Logging helper module
 import logging.config
 import yaml
 
-with open("elise/src/configs/logging_config.yaml", "r", encoding="utf-8") as f:
+with open(
+    "/home/kave/work/Elise/elise/src/configs/logging_config.yaml", "r", encoding="utf-8"
+) as f:
     config = yaml.safe_load(f.read())
     logging.config.dictConfig(config)
     logging.captureWarnings(True)