fix: update imports and dependencies to address vulnerabilities

langtest/datahandler/datasource.py

@@ -1393,7 +1393,7 @@ def __init__(self, dataset: dict, task: TaskManager):
             task (str): Task to be evaluated on.
         """
         self.dataset_name = dataset["data_source"]
-        self.sub_name = dataset.get("subset", "sst2")
+        self.sub_name = dataset.get("subset", "SetFit/sst2")
         self.task = task
 
     @staticmethod

langtest/embeddings/huggingface.py

@@ -13,18 +13,30 @@ class HuggingfaceEmbeddings:
         model (transformers.AutoModel): The transformer model used for sentence embeddings.
     """
 
+    model = None
+    tokenizer = None
+
     def __init__(
         self,
         model: str = "sentence-transformers/all-mpnet-base-v2",
     ):
         """Constructor method
 
         Args:
-            model_name (str): The name of the model to be loaded. By default, it uses the multilingual MiniLM model.
+            model (str): The name of the model to be loaded. By default, it uses the all-mpnet-base-v2 model.
         """
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.tokenizer = AutoTokenizer.from_pretrained(model)
-        self.model = AutoModel.from_pretrained(model).to(self.device)
+
+        # Load model and tokenizer if not already loaded or if model name differs
+        model_loaded = (
+            HuggingfaceEmbeddings.model is not None
+            and HuggingfaceEmbeddings.tokenizer is not None
+            and HuggingfaceEmbeddings.model.config.name_or_path == model
+        )
+
+        if not model_loaded:
+            HuggingfaceEmbeddings.model = AutoModel.from_pretrained(model).to(self.device)
+            HuggingfaceEmbeddings.tokenizer = AutoTokenizer.from_pretrained(model)
 
     def mean_pooling(
         self, model_output: Tuple[torch.Tensor], attention_mask: torch.Tensor

langtest/errors.py

@@ -276,6 +276,8 @@ class Errors(metaclass=ErrorsWithCodes):
     E095 = ("Failed to make API request: {e}")
     E096 = ("Failed to generate the templates in Augmentation: {msg}")
     E097 = ("Failed to load openai. Please install it using `pip install openai`")
+    E098 = ("Invalid model architecture! "
+            "Expected model types are: {model_arch}, but got: {type_model}")
 
 
 class ColumnNameError(Exception):

langtest/modelhandler/transformers_modelhandler.py

@@ -1,8 +1,7 @@
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 import logging
 import numpy as np
 from functools import lru_cache
-from transformers import Pipeline, pipeline, AutoModelForCausalLM, AutoTokenizer
 from .modelhandler import ModelAPI
 from ..utils.custom_types import (
     NEROutput,
@@ -18,6 +17,16 @@
 )
 from ..utils.hf_utils import HuggingFacePipeline
 
+from transformers import (
+    AutoConfig,
+    AutoModelForSeq2SeqLM,
+    Pipeline,
+    pipeline,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    PretrainedConfig,
+)
+
 
 class PretrainedModelForNER(ModelAPI):
     """Transformers pretrained model for NER tasks
@@ -407,62 +416,193 @@ def __call__(
 
 
 class PretrainedModelForTranslation(ModelAPI):
-    """Transformers pretrained model for translation tasks
+    """Transformers pretrained model for translation tasks.
 
     Args:
-        model (transformers.pipeline.Pipeline): Pretrained HuggingFace translation pipeline for predictions.
+        tokenizer: Pretrained HuggingFace tokenizer.
+        model: Pretrained HuggingFace sequence-to-sequence or causal LM model.
+        source_lang (str): Default source language for translation.
+        target_lang (str): Default target language for translation.
+        is_encoder_decoder (bool): Whether the loaded model is an encoder-decoder type.
     """
 
-    def __init__(self, model):
-        """Constructor method
+    def __init__(
+        self,
+        tokenizer: AutoTokenizer,
+        model: Any,
+        source_lang: str = "English",
+        target_lang: str = "German",
+        is_encoder_decoder: bool = True,
+        prompt: Optional[str] = None,
+        **kwargs,
+    ):
+        # Extract the base architecture string names from the model's config
+        model_architectures = (
+            getattr(model, "config", PretrainedConfig()).architectures or []
+        )
 
-        Args:
-            model (transformers.pipeline.Pipeline): Pretrained HuggingFace NER pipeline for predictions.
-        """
-        assert isinstance(model, Pipeline), ValueError(
-            Errors.E079(Pipeline=Pipeline, type_model=type(model))
+        is_supported = any(
+            "CausalLM" in arch
+            or "Seq2SeqLM" in arch
+            or "ConditionalGeneration" in arch
+            or "MTModel" in arch
+            for arch in model_architectures
         )
+
+        assert is_supported, ValueError(
+            Errors.E098(
+                model_arch=(AutoModelForSeq2SeqLM, AutoModelForCausalLM),
+                type_model=type(model),
+            )
+        )
+
+        self.tokenizer = tokenizer
         self.model = model
+        self.source_lang = source_lang
+        self.target_lang = target_lang
+        self.is_encoder_decoder = is_encoder_decoder
+        self.prompt: Optional[str] = prompt
+        self.kwargs = kwargs
 
     @classmethod
-    def load_model(cls, path: str, *args, **kwargs) -> "Pipeline":
-        """Load the Translation model into the `model` attribute.
+    def load_model(cls, path: str, *args, **kwargs) -> "PretrainedModelForTranslation":
+        """Load the Translation model into the `model` attribute."""
 
-        Args:
-            path (str):
-                path to model or model name
-
-        Returns:
-            'Pipeline':
-        """
-        from ..langtest import HARNESS_CONFIG as harness_config
+        source_lang = kwargs.pop("source_language", None)
+        target_lang = kwargs.pop("target_language", None)
+        prompt = kwargs.pop("prompt", None)
 
-        config = harness_config["model_parameters"]
-        tgt_lang = config.get("target_language") or kwargs.get("target_language")
+        # Fallback to langtest configurations if not explicitly provided
+        if not source_lang or not target_lang:
+            try:
+                from langtest import HARNESS_CONFIG
 
-        if "t5" in path:
-            return cls(pipeline(f"translation_en_to_{tgt_lang}", model=path))
+                config_harness = HARNESS_CONFIG.get("model_parameters", {})
+                source_lang = source_lang or config_harness.get(
+                    "source_language", "English"
+                )
+                target_lang = target_lang or config_harness.get(
+                    "target_language", "German"
+                )
+            except ImportError:
+                source_lang = source_lang or "English"
+                target_lang = target_lang or "German"
+
+        tokenizer_loaded = AutoTokenizer.from_pretrained(path)
+        model_config = AutoConfig.from_pretrained(path)
+        is_enc_dec = model_config.is_encoder_decoder
+
+        # Pass remaining args/kwargs (like device_map) to the model loading
+        if is_enc_dec:
+            model_loaded = AutoModelForSeq2SeqLM.from_pretrained(
+                path, device_map="auto", *args, **kwargs
+            )
         else:
-            return cls(pipeline(model=path, src_lang="en", tgt_lang=tgt_lang))
+            print(
+                f"Warning: Model '{path}' is not an encoder-decoder model. "
+                "Loading as AutoModelForCausalLM. Performance may vary."
+            )
+            model_loaded = AutoModelForCausalLM.from_pretrained(
+                path, device_map="auto", *args, **kwargs
+            )
+
+        return cls(
+            tokenizer_loaded,
+            model_loaded,
+            source_lang,
+            target_lang,
+            is_encoder_decoder=is_enc_dec,
+            prompt=prompt,
+        )
 
     @lru_cache(maxsize=102400)
-    def predict(self, text: str, **kwargs) -> TranslationOutput:
-        """Perform predictions on the input text.
+    def predict(
+        self,
+        text: str,
+        **kwargs,
+    ) -> TranslationOutput:
+        """
+        Perform predictions on the input text. Wraps kwargs to enable LRU caching.
 
         Args:
             text (str): Input text to perform translation on.
-            kwargs: Additional keyword arguments.
-
 
         Returns:
-            TranslationOutput: Output model for translation tasks
+            TranslationOutput: Translated text from the input text.
         """
-        prediction = self.model(text, **kwargs)[0]["translation_text"]
+        model_type = getattr(self.model.config, "model_type", "").lower()
+
+        # 1. Handle Prompting for T5 or Custom Prompts
+        if self.prompt:
+            input_text = self.prompt.format(text=text)
+        elif "t5" in model_type:
+            input_text = f"translate {self.source_lang} to {self.target_lang}: {text}"
+        else:
+            input_text = text
+
+        # 2. Set source language for multilingual tokenizers (NLLB, M2M100, mBART)
+        if hasattr(self.tokenizer, "src_lang"):
+            self.tokenizer.src_lang = self.source_lang
+
+        inputs = self.tokenizer(input_text, return_tensors="pt").to(self.model.device)
+
+        # 3. Safely extract and construct generation arguments
+        gen_kwargs = {
+            "max_length": kwargs.pop("max_length", 512),
+            "num_beams": kwargs.pop("num_beams", 1),
+            **kwargs,
+        }
+
+        # Hugging Face crashes if diversity_penalty is set without num_beam_groups > 1
+        num_beam_groups = kwargs.pop("num_beam_groups", 1)
+        if num_beam_groups > 1:
+            gen_kwargs["num_beam_groups"] = num_beam_groups
+            gen_kwargs["diversity_penalty"] = kwargs.pop("diversity_penalty", 0.0)
+
+        # 4. Handle Target Language tokens (forced_bos_token_id)
+        if "forced_bos_token_id" not in gen_kwargs:
+            # Primary strategy: Direct vocabulary lookup (NLLB uses this)
+            if hasattr(self.tokenizer, "convert_tokens_to_ids"):
+                token_id = self.tokenizer.convert_tokens_to_ids(self.target_lang)
+                # Ensure it didn't return an unknown token (meaning target_lang isn't formatted correctly)
+                if token_id is not None and token_id != self.tokenizer.unk_token_id:
+                    gen_kwargs["forced_bos_token_id"] = token_id
+
+            # Fallback strategies for mBART or M2M100
+            if "forced_bos_token_id" not in gen_kwargs:
+                if (
+                    hasattr(self.tokenizer, "lang_code_to_id")
+                    and self.target_lang in self.tokenizer.lang_code_to_id
+                ):
+                    gen_kwargs["forced_bos_token_id"] = self.tokenizer.lang_code_to_id[
+                        self.target_lang
+                    ]
+                elif hasattr(self.tokenizer, "get_lang_id"):
+                    try:
+                        gen_kwargs["forced_bos_token_id"] = self.tokenizer.get_lang_id(
+                            self.target_lang
+                        )
+                    except Exception:
+                        pass
+
+        # Generate tokens
+        output_tokens = self.model.generate(**inputs, **gen_kwargs)
+
+        # Decode based on architecture
+        if self.is_encoder_decoder:
+            prediction = self.tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+        else:
+            # For Causal LMs, strip the input prompt from the output
+            input_length = inputs["input_ids"].shape[1]
+            prediction = self.tokenizer.decode(
+                output_tokens[0][input_length:], skip_special_tokens=True
+            )
+
         return TranslationOutput(translation_text=prediction)
 
     def __call__(self, text: str, *args, **kwargs) -> TranslationOutput:
         """Alias of the 'predict' method"""
-        return self.predict(text=text, **kwargs)
+        return self.predict(text=text, *args, **kwargs)
 
 
 class PretrainedModelForWinoBias(ModelAPI):

langtest/pipelines/transformers/ner_pipeline.py

@@ -126,7 +126,7 @@ def train(self):
             model=self.model,
             args=TrainingArguments(output_dir=self.output_dir, **self.training_args),
             train_dataset=self.train_dataset,
-            tokenizer=self.tokenizer,
+            processing_class=self.tokenizer,
         )
         trainer.train()
         self.model.save_pretrained(self.output_dir)
@@ -215,7 +215,7 @@ def retrain(self):
             ),
             train_dataset=self.augmented_train_dataset,
             eval_dataset=self.eval_dataset,
-            tokenizer=self.tokenizer,
+            processing_class=self.tokenizer,
         )
         trainer.train()
         self.model.save_pretrained(f"augmented_{self.output_dir}")

langtest/prompts.py

@@ -32,7 +32,8 @@ def get_template(self):
         order_less = []
 
         sorted_fields = sorted(
-            self.__dict__.keys(), key=lambda x: self.__field_order.index(x.lower())
+            (k for k in self.__dict__.keys() if k.lower() in self.__field_order),
+            key=lambda x: self.__field_order.index(x.lower()),
         )
 
         for field in sorted_fields:
@@ -126,7 +127,7 @@ def get_input_variables(self):
     def prompt_style(self):
         """Generate a prompt based on the prompt type."""
         if self.prompt_type in ["chat", "instruct"]:
-            from langchain.prompts import (
+            from langchain_core.prompts import (
                 ChatPromptTemplate,
                 FewShotChatMessagePromptTemplate,
             )
@@ -149,7 +150,7 @@ def prompt_style(self):
             return final_prompt
 
         elif self.prompt_type == "completion":
-            from langchain.prompts import FewShotPromptTemplate, PromptTemplate
+            from langchain_core.prompts import FewShotPromptTemplate, PromptTemplate
 
             template = "".join(v for _, v in self.get_template)
             template = f"{template.replace('Answer:', '')}"

langtest/transform/constants.py

@@ -146203,10 +146203,10 @@
 bad_word_list= ['2g1c', '2 girls 1 cup', 'acrotomophilia', 'alabama hot pocket', 'alaskan pipeline', 'anal', 'anilingus', 'anus', 'apeshit', 'arsehole', 'ass', 'asshole', 'assmunch', 'auto erotic', 'autoerotic', 'babeland', 'baby batter', 'baby juice', 'ball gag', 'ball gravy', 'ball kicking', 'ball licking', 'ball sack', 'ball sucking', 'bangbros', 'bangbus', 'bareback', 'barely legal', 'barenaked', 'bastard', 'bastardo', 'bastinado', 'bbw', 'bdsm', 'beaner', 'beaners', 'beaver cleaver', 'beaver lips', 'beastiality', 'bestiality', 'big black', 'big breasts', 'big knockers', 'big tits', 'bimbos', 'birdlock', 'bitch', 'bitches', 'black cock', 'blonde action', 'blonde on blonde action', 'blowjob', 'blow job', 'blow your load', 'blue waffle', 'blumpkin', 'bollocks', 'bondage', 'boner', 'boob', 'boobs', 'booty call', 'brown showers', 'brunette action', 'bukkake', 'bulldyke', 'bullet vibe', 'bullshit', 'bung hole', 'bunghole', 'busty', 'butt', 'buttcheeks', 'butthole', 'camel toe', 'camgirl', 'camslut', 'camwhore', 'carpet muncher', 'carpetmuncher', 'chocolate rosebuds', 'cialis', 'circlejerk', 'cleveland steamer', 'clit', 'clitoris', 'clover clamps', 'clusterfuck', 'cock', 'cocks', 'coprolagnia', 'coprophilia', 'cornhole', 'coon', 'coons', 'creampie', 'cum', 'cumming', 'cumshot', 'cumshots', 'cunnilingus', 'cunt', 'darkie', 'date rape', 'daterape', 'deep throat', 'deepthroat', 'dendrophilia', 'dick', 'dildo', 'dingleberry', 'dingleberries', 'dirty pillows', 'dirty sanchez', 'doggie style', 'doggiestyle', 'doggy style', 'doggystyle', 'dog style', 'dolcett', 'domination', 'dominatrix', 'dommes', 'donkey punch', 'double dong', 'double penetration', 'dp action', 'dry hump', 'dvda', 'eat my ass', 'ecchi', 'ejaculation', 'erotic', 'erotism', 'escort', 'eunuch', 'fag', 'faggot', 'fecal', 'felch', 'fellatio', 'feltch', 'female squirting', 'femdom', 'figging', 'fingerbang', 'fingering', 'fisting', 'foot fetish', 'footjob', 'frotting', 'fuck', 'fuck buttons', 'fuckin', 'fucking', 'fucktards', 'fudge packer', 'fudgepacker', 'futanari', 'gangbang', 'gang bang', 'gay sex', 'genitals', 'giant cock', 'girl on', 'girl on top', 'girls gone wild', 'goatcx', 'goatse', 'god damn', 'gokkun', 'golden shower', 'goodpoop', 'goo girl', 'goregasm', 'grope', 'group sex', 'g-spot', 'guro', 'hand job', 'handjob', 'hard core', 'hardcore', 'hentai', 'homoerotic', 'honkey', 'hooker', 'horny', 'hot carl', 'hot chick', 'how to kill', 'how to murder', 'huge fat', 'humping', 'incest', 'intercourse', 'jack off', 'jail bait', 'jailbait', 'jelly donut', 'jerk off', 'jigaboo', 'jiggaboo', 'jiggerboo', 'jizz', 'juggs', 'kike', 'kinbaku', 'kinkster', 'kinky', 'knobbing', 'leather restraint', 'leather straight jacket', 'lemon party', 'livesex', 'lolita', 'lovemaking', 'make me come', 'male squirting', 'masturbate', 'masturbating', 'masturbation', 'menage a trois', 'milf', 'missionary position', 'mong', 'motherfucker', 'mound of venus', 'mr hands', 'muff diver', 'muffdiving', 'nambla', 'nawashi', 'negro', 'neonazi', 'nigga', 'nigger', 'nig nog', 'nimphomania', 'nipple', 'nipples', 'nsfw', 'nsfw images', 'nude', 'nudity', 'nutten', 'nympho', 'nymphomania', 'octopussy', 'omorashi', 'one cup two girls', 'one guy one jar', 'orgasm', 'orgy', 'paedophile', 'paki', 'panties', 'panty', 'pedobear', 'pedophile', 'pegging', 'penis', 'phone sex', 'piece of shit', 'pikey', 'pissing', 'piss pig', 'pisspig', 'playboy', 'pleasure chest', 'pole smoker', 'ponyplay', 'poof', 'poon', 'poontang', 'punany', 'poop chute', 'poopchute', 'porn', 'porno', 'pornography', 'prince albert piercing', 'pthc', 'pubes', 'pussy', 'queaf', 'queef', 'quim', 'raghead', 'raging boner', 'rape', 'raping', 'rapist', 'rectum', 'reverse cowgirl', 'rimjob', 'rimming', 'rosy palm', 'rosy palm and her 5 sisters', 'rusty trombone', 'sadism', 'santorum', 'scat', 'schlong', 'scissoring', 'semen', 'sex', 'sexcam', 'sexo', 'sexy', 'sexual', 'sexually', 'sexuality', 'shaved beaver', 'shaved pussy', 'shemale', 'shibari', 'shit', 'shitblimp', 'shitty', 'shota', 'shrimping', 'skeet', 'slanteye', 'slut', 's&m', 'smut', 'snatch', 'snowballing', 'sodomize', 'sodomy', 'spastic', 'spic', 'splooge', 'splooge moose', 'spooge', 'spread legs', 'spunk', 'strap on', 'strapon', 'strappado', 'strip club', 'style doggy', 'suck', 'sucks', 'suicide girls', 'sultry women', 'swastika', 'swinger', 'tainted love', 'taste my', 'tea bagging', 'threesome', 'throating', 'thumbzilla', 'tied up', 'tight white', 'tit', 'tits', 'titties', 'titty', 'tongue in a', 'topless', 'tosser', 'towelhead', 'tranny', 'tribadism', 'tub girl', 'tubgirl', 'tushy', 'twat', 'twink', 'twinkie', 'two girls one cup', 'undressing', 'upskirt', 'urethra play', 'urophilia', 'vagina', 'venus mound', 'viagra', 'vibrator', 'violet wand', 'vorarephilia', 'voyeur', 'voyeurweb', 'voyuer', 'vulva', 'wank', 'wetback', 'wet dream', 'white power', 'whore', 'worldsex', 'wrapping men', 'wrinkled starfish', 'xx', 'xxx', 'yaoi', 'yellow showers', 'yiffy', 'zoophilia', '🖕']
 
 DATASETS = {
-    'sst2': {
+    'SetFit/sst2': {
         'subsets': [''],
         'labels': [['Negative Sentiment', 'Positive Sentiment']],
-        'fields': [['sentence']],
+        'fields': [['text']],
         'label_names': ['label'],
     },
     'rotten_tomatoes': {