fix test, add analyse

Author: Wonters

analyse.ipynb

@@ -22,6 +22,7 @@ class TorchModelTrainMixin:
     checkpoint: str = ""
     lr: float = 2e-5
     device: torch.device
+    format_labels_as_long: bool = True
 
     def sample_dataset(self, frac=0.1):
         dataset_size = len(self.dataset)
@@ -53,10 +54,11 @@ def params_optim(self, trial):
 
     def objective(self, trial):
         kwargs = self.params_optim(trial)
-        with mlflow.start_run(nested=True):
-            mlflow.log_params(kwargs)
-            self.reinit_scheduler_optimizer(**kwargs)
-            acc = self.train()
+        if self.tracking:
+            with mlflow.start_run(nested=True):
+                mlflow.log_params(kwargs)
+                self.reinit_scheduler_optimizer(**kwargs)
+        acc = self.train()
         return acc
 
     def get_ddp_dataloader(self, frac=1.0):
@@ -71,7 +73,7 @@ def _train_batch(self, x, y):
         if isinstance(inputs, dict) and inputs["input_ids"]:
             inputs["input_ids"] = inputs["input_ids"].float()
         # todo : fix this case for bert vs lstm
-        if True and isinstance(y, torch.Tensor) and y.dtype == torch.float32:
+        if self.format_labels_as_long and isinstance(y, torch.Tensor) and y.dtype == torch.float32:
             labels = y.long()
         else:
             labels = y.float()
@@ -100,10 +102,12 @@ def _train_batch(self, x, y):
             else:
                 raise e
         self.optimizer.step()
-        logger.info(f" Rank {dist.get_rank()} loss {loss.item()} acc {acc}")
-        mlflow.log_metric("loss", loss.item())
-        mlflow.log_metric("acc", acc)
-        mlflow.log_metric("time", time.time())
+        if dist.is_initialized():
+            logger.info(f" Rank {dist.get_rank()} loss {loss.item()} acc {acc}")
+        if self.tracking:
+            mlflow.log_metric("loss", loss.item())
+            mlflow.log_metric("acc", acc)
+            mlflow.log_metric("time", time.time())
         del inputs, labels, outputs, loss
         gc.collect()
         if torch.backends.mps.is_available():

src/mixins.py

@@ -33,6 +33,8 @@
 import mlflow
 from mlflow.data.pandas_dataset import from_pandas
 from mlflow.models import infer_signature
+from transformers.models.tapas.modeling_tapas import flatten
+
 from .mixins import TorchModelTrainMixin
 from .torch_models import LSTMTorchNN
 from .dataset import TweetDataset
@@ -72,11 +74,11 @@ class BaseModelABC(ABC):
     dataset_class = None
     tokenizer_class = None
     batch_size = 32
-    #artifact_uri = "file:///app/mlruns"
 
-    def __init__(self, dataset: pd.DataFrame):
+    def __init__(self, dataset: pd.DataFrame, tracking: bool = True):
+        self.tracking = tracking
         self.dataset = None
-        if dataset:
+        if dataset is not None:
             self.original_dataset = dataset
             self.x_train, self.x_test, self.x_val, self.y_train, self.y_test, self.y_val = split_data(dataset)
             self.name = self.__class__.__name__
@@ -112,10 +114,12 @@ def init_items(self):
         """
         self.model = None
         self.tokenizer = None
+        self.tokenizer.fit_transform(self.x_train)
 
     def init_mlflow(self, name:str = ""):
-        self.run = mlflow.start_run(run_name=name if name else self.name)
-        self.run_id = self.run.info.run_id
+        if self.tracking:
+            self.run = mlflow.start_run(run_name=name if name else self.name)
+            self.run_id = self.run.info.run_id
 
     def load_checkpoint(self) -> object:
         """
@@ -135,11 +139,13 @@ def preprocessing(self, data):
 
     def confusion_matrix(self):
         with NamedTemporaryFile(suffix=".png") as f:
-            conf_mat = confusion_matrix(self.y_train, self.predict(list(self.x_train)))
+            y_pred = self.predict(list(self.x_train), flatten=True)
+            conf_mat = confusion_matrix(self.y_train, y_pred)
             group_names = ["True Neg", "False Pos", "False Neg", "True Pos"]
             group_counts = [f"{value: 0.0f}" for value in conf_mat.flatten()]
             group_percentages = [f"{value:.2%}" for value in conf_mat.flatten() / np.sum(conf_mat)]
             labels = [f"{v1}\n{v2}\n{v3}" for v1, v2, v3 in zip(group_names, group_counts, group_percentages)]
+            print(labels)
             labels = np.asarray(labels).reshape(2, 2)
             sns.heatmap(
                 conf_mat,
@@ -151,46 +157,49 @@ def confusion_matrix(self):
             plt.ylabel("Cluster réels")
             plt.savefig(f.name)
             plt.close()
-            mlflow.log_artifact(f.name, "confusion_matrix.png")
+            if self.tracking:
+                mlflow.log_artifact(f.name, "confusion_matrix.png")
 
     def train(self):
         """
         Train the model here
         """
         self.save()
-        mlflow.set_tag("model_type", self.name)
-        self.log_metrics()
+        if self.tracking:
+            mlflow.set_tag("model_type", self.name)
+            self.log_metrics()
         self.confusion_matrix()
-        mlflow.log_artifact(self.checkpoint)
-        signature = infer_signature(self.x_train, self.predict(self.x_train))
-        dataset = from_pandas(self.original_dataset.loc[self.x_train.index], source="local")
-        mlflow.log_input(dataset, context="tweet-dataset")
-        if isinstance(self.model, PreTrainedModel):
-            mlflow.transformers.log_model(
-                transformers_model=self.checkpoint,
-                artifact_path=self.name,
-                task="text-classification",  # important !
-                tokenizer=self.tokenizer,
-                signature=signature,
-                registered_model_name=f"{self.name}-quickstart"
-            )
-        elif isinstance(self.model, lgm.LGBMClassifier):
-            mlflow.lightgbm.log_model(
-                lgb_model=self.model,
-                artifact_path=self.name,
-                signature=signature,
-                registered_model_name=f"{self.name}-quickstart",
-            )
-        else:
-            mlflow.sklearn.log_model(
-                sk_model=self.model,
-                artifact_path=self.name,
-                signature=signature,
-                registered_model_name=f"{self.name}-quickstart",
-            )
-        mlflow.end_run()
+        if self.tracking:
+            mlflow.log_artifact(self.checkpoint)
+            signature = infer_signature(self.x_train, self.predict(self.x_train))
+            dataset = from_pandas(self.original_dataset.loc[self.x_train.index], source="local")
+            mlflow.log_input(dataset, context="tweet-dataset")
+            if isinstance(self.model, PreTrainedModel):
+                mlflow.transformers.log_model(
+                    transformers_model=self.checkpoint,
+                    artifact_path=self.name,
+                    task="text-classification",  # important !
+                    tokenizer=self.tokenizer,
+                    signature=signature,
+                    registered_model_name=f"{self.name}-quickstart"
+                )
+            elif isinstance(self.model, lgm.LGBMClassifier):
+                mlflow.lightgbm.log_model(
+                    lgb_model=self.model,
+                    artifact_path=self.name,
+                    signature=signature,
+                    registered_model_name=f"{self.name}-quickstart",
+                )
+            else:
+                mlflow.sklearn.log_model(
+                    sk_model=self.model,
+                    artifact_path=self.name,
+                    signature=signature,
+                    registered_model_name=f"{self.name}-quickstart",
+                )
+            mlflow.end_run()
 
-    def predict(self, x: Union[pd.Series, np.ndarray]):
+    def predict(self, x: Union[pd.Series, np.ndarray], flatten: bool = True):
         """
         Predict the sentiment of the input data
         """
@@ -200,7 +209,6 @@ def predict(self, x: Union[pd.Series, np.ndarray]):
 class SklearnBaseModel(BaseModelABC):
     def log_metrics(self):
         super().log_metrics()
-        #mlflow.sklearn.log_model(self.model, self.name)
         mlflow.log_params(self.model.get_params())
 
 class TorchBaseModel(TorchModelTrainMixin, BaseModelABC):
@@ -209,13 +217,13 @@ class TorchBaseModel(TorchModelTrainMixin, BaseModelABC):
     """
     distributed = False
 
-    def __init__(self, dataset: pd.DataFrame):
+    def __init__(self, dataset: pd.DataFrame, tracking: bool = True):
         if self.distributed:
             dist.init_process_group("nccl")
         if dist.is_initialized():
             self.local_rank = dist.get_rank()
             torch.cuda.set_device(self.local_rank)
-        super().__init__(dataset)
+        super().__init__(dataset, tracking)
         if dist.is_initialized():
             self.dataloader, self.sampler = self.get_ddp_dataloader()
             logger.info(f"Rank {dist.get_rank()} using DDP")
@@ -257,6 +265,8 @@ def init_items(self):
         self.tokenizer = self.tokenizer_class(
             max_features=1000, ngram_range=(1, 2), binary=True
         )
+        self.tokenizer.fit_transform(self.x_train)
+
 
     def train(self):
         """
@@ -297,6 +307,7 @@ def init_items(self):
         self.tokenizer = self.tokenizer_class(
             max_features=1000, min_df=2, max_df=0.95
         )
+        self.tokenizer.fit_transform(self.x_train)
 
     def clean(self, tweet):
         translator = str.maketrans('','', string.punctuation)
@@ -348,10 +359,11 @@ def init_items(self):
         """
         self.model = LogisticRegression(max_iter=1000,
                                         C=1.7279373898388395,
-                                        penalty='l1',
+                                        penalty='l2',
                                         n_jobs=4, 
                                         verbose=True)
         self.tokenizer = self.tokenizer_class()
+        self.tokenizer.fit_transform(self.x_train)
 
     def objective(self, tokens, params):
         with mlflow.start_run(nested=True):
@@ -431,6 +443,7 @@ def load_checkpoint(self):
                 num_training_steps=total_steps
             )
             self.criterion = torch.nn.CrossEntropyLoss()
+        self.model.to(self.device)
 
     def save(self):
         if self.distributed and dist.is_initialized():
@@ -439,7 +452,7 @@ def save(self):
             self.model.save_pretrained(self.checkpoint)
         self.tokenizer.save_pretrained(self.checkpoint)
 
-    def predict(self, x: list):
+    def predict(self, x: list, flatten: bool = False):
         predicted_class = []
         for i in range(0, len(x), self.batch_size):
             inputs = self.preprocessing(x[i:i+self.batch_size])
@@ -448,10 +461,13 @@ def predict(self, x: list):
                 outputs = self.model(**inputs)
                 probs = F.softmax(outputs.logits, dim=1)
                 confidence, categorie = probs.max(dim=1)
-                predicted_class.extend([{'prediction': categorie.item(), 
-                                         'confidence': confidence.item()} 
-                                         for confidence, categorie in zip(confidence, categorie)
-                                         ])
+                if not flatten:
+                    predicted_class.extend([{'prediction': categorie.item(),
+                                             'confidence': confidence.item()}
+                                             for confidence, categorie in zip(confidence, categorie)
+                                             ])
+                else:
+                    predicted_class.extend([c.item() for c in categorie])
         return predicted_class
 
 class RobertaModel(BertModel):
@@ -474,7 +490,7 @@ class LSTMModel(TorchBaseModel):
     out_features = 1
     lr = 1e-4
     device = DEVICE
-    # torch.nn.CrossEntropyLoss()
+    format_labels_as_long: bool = False
 
     @property
     def get_metrics(self) -> dict:
@@ -540,7 +556,7 @@ def save(self):
         self.tokenizer.save_pretrained(self.checkpoint)
         torch.save(self.model.state_dict(), self.checkpoint + "/model.pth")
 
-    def predict(self, x):
+    def predict(self, x, flatten: bool = True):
         predicted_class = []
         for i in range(0, len(x), self.batch_size):
             inputs = self.preprocessing(x[i:i+self.batch_size])

src/ml.py

@@ -4,9 +4,9 @@
 from fastapi.routing import APIRouter
 from fastapi.requests import Request
 from fastapi.responses import Response
-from fastapi import Form, WebSocket
+from fastapi import WebSocket
 from fastapi.templating import Jinja2Templates
-from multiprocessing import Process, Pipe
+from multiprocessing import Pipe
 from multiprocessing.pool import Pool
 from typing import List
 from rich.logging import RichHandler
@@ -15,10 +15,9 @@
 import pymongo
 import uuid
 import time
-import random
 import json
-from .ml import BertModel, RobertaModel
-from .models import Tweet, Sentiment
+from .ml import RobertaModel
+from .models import Tweet
 
 # Configuration des métriques Prometheus
 PREDICTION_COUNT = prom.Counter(
@@ -55,13 +54,21 @@ def get_pool():
     return pool
 
 
-def run_predict(text: List[Tweet], sender):
+def run_predict(text: List[Tweet], sender, save_db=False):
     FLAG_START = "started"
     FLAG_DONE = "done"
     logger.info(f"prediction started {text}")
     sender.send(FLAG_START)
     start_time = time.time()
     result = RobertaModel(dataset=None).predict([t.text for t in text])
+    if save_db:
+        with pymongo.MongoClient(MONGO_URI) as client:
+            db = client["sentiment_analyses"]
+            collection = db["tweets"]
+            # Convertir la liste de tweets en liste de documents
+            tweets = [{"text": str(t), "prediction": r['prediction'], "confidence": r['confidence']} for r, t in zip(result, text)]
+            collection.insert_many(tweets)
+            logger.info(f"tweets added to db: {len(tweets)} tweets")
     logger.info(f"prediction done {text}")
     sender.send(FLAG_DONE)
     result = [{'prediction': r['prediction'], 'confidence': r['confidence'], 'text': t.text} for r, t in zip(result, text)]
@@ -70,9 +77,8 @@ def run_predict(text: List[Tweet], sender):
 
 class PredictApp:
     ACK_TIMEOUT = 1.0
-    
 
-    def __init__(self):
+    def __init__(self, save_db=False):
         self.router = APIRouter()
         self.router.add_api_route("/", self.get, methods=["GET"])
         self.router.add_api_route("/predict", self.predict, methods=["POST"])
@@ -81,6 +87,7 @@ def __init__(self):
         self.active_connections = {}
         self.tasks = {}
         self.pipes = {}
+        self.save_db = save_db
 
     async def get(self, request: Request):
         """"""
@@ -178,16 +185,9 @@ async def predict(self, request: Request, text: List[Tweet]):
         """
         Predict
         """
-        with pymongo.MongoClient(MONGO_URI) as client:
-            db = client["sentiment_analyses"]
-            collection = db["tweets"]
-            # Convertir la liste de tweets en liste de documents
-            tweets = [{"text": str(tweet)} for tweet in text]
-            collection.insert_many(tweets)
-            logger.info(f"tweets added to db: {len(tweets)} tweets")
         p = get_pool()
         pipe = Pipe()
-        result = p.apply_async(run_predict, (text, pipe[1]))
+        result = p.apply_async(run_predict, (text, pipe[1], self.save_db))
         PREDICTION_STATUS.labels("bert").inc()
         task_id = str(uuid.uuid4())
         self.tasks[task_id] = result
@@ -196,6 +196,6 @@ async def predict(self, request: Request, text: List[Tweet]):
         return {"task_id": task_id, "status": "processing"}
 
 
-predict_app = PredictApp()
+predict_app = PredictApp(save_db=True)
 app.include_router(predict_app.router)