lmspell 1.0.0__py3-none-any.whl

lmspell/__init__.py

@@ -0,0 +1,47 @@
+"""LMSpell: spell correction with pre-trained language models.
+
+Heavy dependencies (torch, transformers, unsloth) are imported lazily, so
+``import lmspell`` stays cheap; they load on first use of ``LMSpell`` or
+``ModelEnum``.
+"""
+
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+from lmspell.exceptions import (
+    EvaluateModelException,
+    InvalidModelException,
+    LMSpellException,
+)
+
+__version__ = "1.0.0"
+
+if TYPE_CHECKING:
+    from lmspell.core import LMSpell
+    from lmspell.model_enum import ModelEnum
+
+__all__ = [
+    "LMSpell",
+    "ModelEnum",
+    "LMSpellException",
+    "InvalidModelException",
+    "EvaluateModelException",
+    "__version__",
+]
+
+_LAZY_ATTRS = {
+    "LMSpell": "lmspell.core",
+    "ModelEnum": "lmspell.model_enum",
+}
+
+
+def __getattr__(name):
+    if name in _LAZY_ATTRS:
+        value = getattr(import_module(_LAZY_ATTRS[name]), name)
+        globals()[name] = value
+        return value
+    raise AttributeError(f"module 'lmspell' has no attribute {name!r}")
+
+
+def __dir__():
+    return sorted(set(globals()) | set(__all__))

lmspell/core.py

@@ -0,0 +1,26 @@
+from lmspell.exceptions import InvalidModelException
+from lmspell.model_enum import ModelEnum
+
+
+class LMSpell:
+
+    def __init__(self, model, language="Sinhala"):
+        """
+        Create a spell corrector for a supported Hugging Face model id (see ModelEnum).
+        `language` is the language name used in the LLM prompt and for the mBART language code.
+        """
+
+        try:
+            model_enum = ModelEnum(model)
+        except ValueError:
+            raise InvalidModelException(f"Invalid model: {model}") from None
+
+        model_class = model_enum.getModelClass()
+        self._model = model_class(model_label=model_enum, language=language)
+
+    def __getattr__(self, name):
+        # Only called for attributes LMSpell does not have; guard against recursion
+        # when _model itself is missing (e.g. during unpickling or a failed __init__).
+        if name == "_model":
+            raise AttributeError(name)
+        return getattr(self._model, name)

lmspell/data/lm_data_loader.py

@@ -0,0 +1,43 @@
+from torch.utils.data import DataLoader
+
+from lmspell.data.lm_spell_dataset import LMSpellDataset
+from lmspell.utils.dataset_utils import DatasetUtils
+
+
+class LMDataLoader:
+
+    def __init__(
+            self,
+            tokenizer,
+            dataset = LMSpellDataset,
+            train_path = None,
+            val_path = None,
+            test_path = None,
+            dataset_size = 1,
+            train_max_length = 128,
+            test_max_length = 128,
+            train_batch_size = 16,
+            test_batch_size = 16,
+        ):
+
+        self.dataset = dataset
+        self.train_dataloader = None
+        self.val_dataloader = None
+        self.test_dataloader = None
+
+        dataset = DatasetUtils(train_path=train_path, val_path=val_path, test_path=test_path, dataset_size=dataset_size)
+
+        if dataset.train_dataset is not None:
+            train_dataset = self.dataset(dataset.train_dataset, tokenizer, train_max_length)
+            train_dataloader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=False, pin_memory=True)
+            self.train_dataloader = train_dataloader
+
+        if dataset.val_dataset is not None:
+            val_dataset = self.dataset(dataset.val_dataset, tokenizer, train_max_length)
+            val_dataloader = DataLoader(val_dataset, batch_size=train_batch_size, shuffle=False, pin_memory=True)
+            self.val_dataloader = val_dataloader
+
+        if dataset.test_dataset is not None:
+            test_dataset = self.dataset(dataset.test_dataset, tokenizer, test_max_length)
+            test_dataloader = DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False, pin_memory=True)
+            self.test_dataloader = test_dataloader

lmspell/data/lm_spell_dataset.py

@@ -0,0 +1,27 @@
+import pandas as pd
+from torch.utils.data import Dataset
+
+from lmspell.utils.config_utils import ConfigUtils
+
+
+class LMSpellDataset(Dataset):
+    def __init__(self, dataframe, tokenizer, max_length=128):
+        if not isinstance(dataframe, pd.DataFrame):
+            dataframe = pd.DataFrame(dataframe)
+        self.dataframe = dataframe
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+    def __len__(self):
+        return len(self.dataframe)
+
+    def __getitem__(self, idx):
+        item = self.dataframe.iloc[idx]
+        dataset_columns = ConfigUtils.get_dataset_columns()
+        inputs = self.tokenizer(item[dataset_columns[0]], return_tensors='pt', padding='max_length', truncation=True, max_length=self.max_length)
+        targets = self.tokenizer(item[dataset_columns[1]], return_tensors='pt', padding='max_length', truncation=True, max_length=self.max_length)
+        inputs['labels'] = targets['input_ids']
+        return {key: val.squeeze() for key, val in inputs.items()}
+
+    def get(self):
+        return len(self.dataframe)

lmspell/exceptions.py

@@ -0,0 +1,13 @@
+"""Exceptions raised by LMSpell."""
+
+
+class LMSpellException(Exception):
+    """Base class for all LMSpell errors."""
+
+
+class InvalidModelException(LMSpellException):
+    """Raised when an unsupported model name is requested."""
+
+
+class EvaluateModelException(LMSpellException):
+    """Raised when evaluation of a model fails."""

lmspell/model_enum.py

@@ -0,0 +1,121 @@
+from enum import Enum
+from typing import Dict, Type
+
+from lmspell.models.model_abstract import ModelAbstract
+
+
+class ModelEnum(Enum):
+    """
+    Enum class for all the models supported by the 
+    Neural Spell Checker Library
+    """
+    # Fine-tuned models
+    # These checkpoints will be released upon the acceptance of the paper
+
+    # Llama 3.1 models
+    LLAMA_31_8B = "unsloth/Meta-Llama-3.1-8B-bnb-4bit"
+    LLAMA_31_8B_INSTRUCT = "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"
+
+    # Llama 3.2 models
+    LLAMA_32_1B = "unsloth/Llama-3.2-1B-bnb-4bit"
+    LLAMA_32_1B_INSTRUCT = "unsloth/Llama-3.2-1B-Instruct-bnb-4bit"
+    LLAMA_32_3B = "unsloth/Llama-3.2-3B-bnb-4bit"
+    LLAMA_32_3B_INSTRUCT = "unsloth/Llama-3.2-3B-Instruct-bnb-4bit"
+
+    # Gemma 2 models
+    GEMMA_2_2B = "unsloth/gemma-2-2b-bnb-4bit"
+    GEMMA_2_2B_INSTRUCT = "unsloth/gemma-2-2b-it-bnb-4bit"
+    GEMMA_2_9B = "unsloth/gemma-2-9b-bnb-4bit"
+    GEMMA_2_9B_INSTRUCT = "unsloth/gemma-2-9b-it-bnb-4bit"
+
+    # MT5 models
+    MT5_SMALL = "google/mt5-small"
+    MT5_BASE = "google/mt5-base"
+    MT5_LARGE = "google/mt5-large"
+    MT5_XL = "google/mt5-xl"
+    MT5_XXL = "google/mt5-xxl"
+
+    # MBART models
+    MBART50 = "facebook/mbart-large-50"
+
+    # SinBERT models
+    SINBERT_SMALL = "NLPC-UOM/SinBERT-small"
+    SINBERT_LARGE = "NLPC-UOM/SinBERT-large"
+
+    # XLMR models
+    XLMR_BASE = "FacebookAI/xlm-roberta-base"
+    XLMR_LARGE = "FacebookAI/xlm-roberta-large"
+
+    def getModelClass(self) -> Type[ModelAbstract]:
+        """Return the model class associated with the enum value"""
+
+        from lmspell.models.gemma.gemma2_9b_model import Gemma29BModel
+        from lmspell.models.gemma.gemma_abstract_model import GemmaAbstractModel
+        from lmspell.models.llama.llama31_8b_model import Llama318BModel
+        from lmspell.models.llama.llama_abstract_model import LlamaAbstractModel
+        from lmspell.models.mbart.mbart50_model import Mbart50Model
+        from lmspell.models.mt5.mt5_abstract_model import Mt5AbstractModel
+        from lmspell.models.mt5.mt5_base_model import Mt5BaseModel
+        from lmspell.models.mt5.mt5_large_model import Mt5LargeModel
+        from lmspell.models.sinbert.sinbert_abstract_model import SinBertAbstractModel
+        from lmspell.models.xlmr.xlmr_abstract_model import XlmrAbstractModel
+        from lmspell.models.xlmr.xlmr_large_model import XlmrLargeModel
+
+        # import other model classes as needed
+
+        model_classes: Dict[ModelEnum, Type[ModelAbstract]] = {
+            # Fine-tuned models
+            # ModelEnum.LLAMA31_FT_SI: LlamaAbstractModel,
+            # ModelEnum.GEMMA2_FT_SI: GemmaAbstractModel,
+            # ModelEnum.MT5_FT_SI: Mt5AbstractModel,
+            # ModelEnum.MT5_XL_FT_SI: Mt5AbstractModel,
+            # ModelEnum.MBART_FT_SI: MbartAbstractModel,
+            # ModelEnum.SINBERT_FT_SI: SinBertAbstractModel,
+            # ModelEnum.XLMR_FT_SI: XlmrAbstractModel,
+
+            # Llama 3.1 models
+            ModelEnum.LLAMA_31_8B: Llama318BModel,
+            ModelEnum.LLAMA_31_8B_INSTRUCT: LlamaAbstractModel,
+
+            # Llama 3.2 models
+            ModelEnum.LLAMA_32_1B: LlamaAbstractModel,
+            ModelEnum.LLAMA_32_1B_INSTRUCT: LlamaAbstractModel,
+            ModelEnum.LLAMA_32_3B: LlamaAbstractModel,
+            ModelEnum.LLAMA_32_3B_INSTRUCT: LlamaAbstractModel,
+
+            # Gemma 2 models
+            ModelEnum.GEMMA_2_2B: GemmaAbstractModel,
+            ModelEnum.GEMMA_2_2B_INSTRUCT: GemmaAbstractModel,
+            ModelEnum.GEMMA_2_9B: Gemma29BModel,
+            ModelEnum.GEMMA_2_9B_INSTRUCT: GemmaAbstractModel,
+
+            # mT5 models
+            ModelEnum.MT5_SMALL: Mt5AbstractModel,
+            ModelEnum.MT5_BASE: Mt5BaseModel,
+            ModelEnum.MT5_LARGE: Mt5LargeModel,
+            ModelEnum.MT5_XL: Mt5AbstractModel,
+            ModelEnum.MT5_XXL: Mt5AbstractModel,
+
+            # mBART models
+            ModelEnum.MBART50: Mbart50Model,
+
+            # SinBERT models
+            ModelEnum.SINBERT_SMALL: SinBertAbstractModel,
+            ModelEnum.SINBERT_LARGE: SinBertAbstractModel,
+
+            # XLMR models
+            ModelEnum.XLMR_BASE: XlmrAbstractModel,
+            ModelEnum.XLMR_LARGE: XlmrLargeModel,
+
+            # Add other model-class mappings here
+        }
+        return model_classes[self]
+
+    def get_model_path(self) -> str:
+        """Return the model path"""
+        return self.value
+
+    @staticmethod
+    def get_supported_models():
+        """Return a list of all supported model names"""
+        return [model.value for model in ModelEnum]

lmspell/models/decoder_abstract.py

@@ -0,0 +1,143 @@
+import pandas as pd
+
+from lmspell.exceptions import LMSpellException
+from lmspell.models.model_abstract import ModelAbstract
+from lmspell.utils.config_utils import ConfigUtils
+from lmspell.utils.evaluate_utils import EvaluateUtils
+
+
+class DecoderAbstract(ModelAbstract):
+    do_sample: bool
+
+    def __init__(self, model_label=None, language="Sinhala"):
+        import torch
+        from unsloth import FastLanguageModel
+
+        super().__init__(model_label, language)
+
+        # Models mapped straight to LlamaAbstractModel/GemmaAbstractModel do not set do_sample;
+        # use the same default as the concrete Llama318BModel/Gemma29BModel.
+        if not hasattr(self, "do_sample"):
+            self.do_sample = True
+
+        self.max_seq_length = 2048    # WHY!!!
+        is_t4 = torch.cuda.is_available() and torch.cuda.get_device_name(0).startswith('Tesla T4')
+        self.dtype = torch.float16 if is_t4 else torch.bfloat16
+        self.load_in_4bit = True    # Get as a parameter
+        self.attn_implementation = "flash_attention_2"
+
+        model, tokenizer = FastLanguageModel.from_pretrained(
+            model_name = self.model_label.get_model_path(),
+            max_seq_length = self.max_seq_length,
+            dtype = self.dtype,
+            load_in_4bit = self.load_in_4bit,
+            attn_implementation=self.attn_implementation,
+            random_state = self.seed,
+        )
+
+        # Need to check pros and cons of loading it like this for both training and inference.
+        model = FastLanguageModel.get_peft_model(
+            model,
+            r = 8, # 8, 16, 32, 64, 128 suggested
+            target_modules = ["q_proj", "k_proj", "v_proj", "o_proj","gate_proj", "up_proj", "down_proj"],
+            lora_alpha = 16,
+            lora_dropout = 0,
+            bias = "none",
+            use_gradient_checkpointing = "unsloth",
+            random_state = self.seed,
+            use_rslora = False,
+            loftq_config = None,
+        )
+
+        self.model = model
+        self.tokenizer = tokenizer
+
+        # pylint: disable-next=line-too-long
+        self.PROMPT = f"""You are an expert {self.language} spell corrector. Below is a sentence in {self.language} language. It may or may not have a spelling mistake. Give the corrected output in {self.language}.
+
+        ### Text:
+        {{}}
+
+        ### Output:
+        {{}}"""
+
+    def correct(self, input_set: list[str] | str | pd.DataFrame, target_set: list[str] | str | pd.DataFrame = None, output_dir: str = 'outputs'):
+        """
+        Corrects the text. The input set should be a string, a list of strings or a DataFrame.
+        Perform the evaluation if the target set is provided.
+        """
+        import os
+        import re
+
+        from tqdm import tqdm
+        from unsloth import FastLanguageModel
+
+        os.makedirs(os.path.join(os.getcwd(), output_dir), exist_ok=True)
+
+        input_set, evaluate_flag = self.process_input(input_set, target_set)
+
+        dataset_col_names = ConfigUtils.get_dataset_columns()
+
+        FastLanguageModel.for_inference(self.model)
+        pred_texts = []
+        generation_mode = self.model.generation_config.get_generation_mode()
+        print(f"Generation mode: {generation_mode}")
+
+        for src in tqdm(input_set[dataset_col_names[0]].to_list(), desc="Generating predictions"):
+            inputs = self.tokenizer(
+                [self.PROMPT.format(src, "")],
+                return_tensors="pt"
+            ).to(self.device)
+
+            outputs = self.model.generate(**inputs, max_new_tokens=self.max_seq_length, do_sample=self.do_sample, use_cache=True)
+
+            pred_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+            pred_texts.append(pred_text)
+
+        output_texts = []
+        for pred in pred_texts:
+            match = re.search(r'### Output:\s*([^\n]+)', pred)
+            output_texts.append(match.group(1).strip() if match else None)
+
+        result_col_names = ConfigUtils.get_results_columns()
+
+        results_data = {
+            result_col_names[0]: input_set[dataset_col_names[0]].to_list(),
+            result_col_names[1]: output_texts
+        }
+
+        results_df = pd.DataFrame(results_data)
+
+        if evaluate_flag:
+            results_df[result_col_names[2]] = input_set[dataset_col_names[1]].to_list()
+            print("Evaluating the outputs...")
+            EvaluateUtils.evaluate_from_dataframe(results_df, output_dir)
+
+        return results_df
+
+    def correctFromFile(self, src: str, target: str = None, output_dir: str = "outputs"):
+        """
+        Corrects the text from a file. The file should be in the format of
+        a CSV or TXT file. If the file is a CSV, it should have two columns: 'text' and 'expected'.
+        If the file is a TXT, it should have one line per sentence. If target is provided,
+        it will be used for evaluation. The output will be saved in the specified output directory.
+        """
+
+        if src.endswith('.csv'):
+            df = pd.read_csv(src)
+            input_col, target_col = ConfigUtils.get_dataset_columns()
+            src = df[input_col].tolist()
+            if target_col in df.columns:
+                target = df[target_col].tolist()
+
+        elif src.endswith('.txt'):
+            with open(src, 'r', encoding='utf-8') as file:
+                src = file.readlines()
+            if target is not None:
+                with open(target, 'r', encoding='utf-8') as file:
+                    target = file.readlines()
+        else:
+            raise LMSpellException("Unsupported file format. Only .csv and .txt are supported.") from None
+
+        return self.correct(src, target, output_dir)

lmspell/models/encoder_abstract.py

@@ -0,0 +1,171 @@
+import os
+from typing import Any
+
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from lmspell.data.lm_spell_dataset import LMSpellDataset
+from lmspell.exceptions import LMSpellException
+from lmspell.models.model_abstract import ModelAbstract
+from lmspell.utils.config_utils import ConfigUtils
+from lmspell.utils.evaluate_utils import EvaluateUtils
+from lmspell.utils.general_utils import GeneralUtils
+from lmspell.utils.pre_post_processing_utils import PrePostProcessingUtils
+
+
+class EncoderAbstract(ModelAbstract):
+    accelerator: Any
+
+    def correct(self, input_set: list[str] | str | pd.DataFrame, target_set: list[str] | str | pd.DataFrame = None,
+                output_dir: str = "outputs", max_length: int = 128, batch_size: int = 8,
+                shuffle: bool = False):  # pylint: disable=unused-argument  # kept for API parity with EncoderDecoderAbstract
+        """
+        Corrects the text using encoder-only models (e.g., XLM-R). Supports string, list, or DataFrame as input.
+        If target_set is provided, performs evaluation and saves results.
+        """
+        os.makedirs(os.path.join(os.getcwd(), output_dir), exist_ok=True)
+
+        if max_length > 514:
+            print("Warning: max_length is set to 514, as it is the maximum length for XLMR models.")
+            max_length = 514
+
+        input_set, evaluate_flag = self.process_input(input_set, target_set)
+
+        dataset = LMSpellDataset(input_set, self.tokenizer, max_length)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, pin_memory=True)
+
+        self.model = self.accelerator.prepare_model(self.model, evaluation_mode=True)
+        dataloader = self.accelerator.prepare_data_loader(dataloader)
+        self.model.resize_token_embeddings(len(self.tokenizer))
+        GeneralUtils.clean_memory()
+
+        self.model.eval()
+        originals, predictions, labels = [], [], []
+
+        with tqdm(dataloader, leave=True, disable=not self.accelerator.is_local_main_process) as pbar:
+            for batch in pbar:
+                with torch.no_grad():
+                    outputs = self.model(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"], labels=batch["labels"])
+
+                pred = torch.argmax(outputs.logits, dim=-1)
+                orig = batch["input_ids"]
+                label = batch["labels"]
+
+                gathered_orig, gathered_pred, gathered_label = self.accelerator.gather_for_metrics((orig, pred, label))
+                self.accelerator.wait_for_everyone()
+
+                originals.extend(gathered_orig)
+                predictions.extend(gathered_pred)
+                labels.extend(gathered_label)
+
+        results_df = self.decode(originals, predictions, labels, output_dir)
+
+        if evaluate_flag and self.accelerator.is_main_process:
+            eval_df = results_df.iloc[:, :2].copy()
+            if target_set is None:
+                # DataFrame input that carries its own target column.
+                target_col = input_set[ConfigUtils.get_dataset_columns()[1]]
+            elif isinstance(target_set, pd.DataFrame):
+                target_col = target_set.iloc[:, 0] if target_set.shape[1] > 0 else pd.Series([])
+            elif isinstance(target_set, list):
+                target_col = pd.Series(target_set)
+            elif isinstance(target_set, str):
+                target_col = pd.Series([target_set])
+            else:
+                raise LMSpellException("Cannot evaluate: target_set must be a string, a list of strings or a DataFrame.")
+
+            eval_df[ConfigUtils.get_results_columns()[2]] = target_col.reset_index(drop=True)
+
+            print("Evaluating the outputs...")
+            EvaluateUtils.evaluate_from_dataframe(eval_df, output_dir)
+
+        return results_df
+
+    def decode(self, originals, predictions, labels, output_dir):
+        """
+        Decode a tensor of tensors and save it to a file.
+        """
+
+        skip_special_tokens = False
+
+        # This is different in mbart model. Same for xlmr, sinbert and mt5. We don't need to do this in large models.
+        special_token_id_to_keep = self.tokenizer.convert_tokens_to_ids(PrePostProcessingUtils.ZWJ)
+        all_special_ids = torch.tensor(self.tokenizer.all_special_ids, dtype=torch.int64).to(self.device)
+
+        special_token_tensor = torch.tensor([special_token_id_to_keep], dtype=torch.int64).to(self.device)
+
+        originals, predictions, labels = [
+            [PrePostProcessingUtils.remove_special_tokens(tokens.to(self.device), all_special_ids, special_token_tensor) for tokens in data]
+            for data in (originals, predictions, labels)
+        ]
+
+        originals_decoded, predictions_decoded, labels_decoded = [
+            [self.tokenizer.decode(ids, skip_special_tokens=skip_special_tokens).replace('\n', '').strip() for ids in data]
+            for data in (originals, predictions, labels)
+        ]
+
+        result_col_names = ConfigUtils.get_results_columns()
+
+        results_data = {
+            result_col_names[0]: originals_decoded,
+            result_col_names[1]: predictions_decoded,
+            result_col_names[2]: labels_decoded,
+        }
+
+        # Write CSV inside the main process condition
+        results_df = pd.DataFrame(results_data)
+
+        # When using BERT models, these character were present in the output. They throws errors when saving in excel format
+        results_df = results_df.apply(lambda col: col.map(lambda x: x.replace('.\uffff', '').replace('\x11', '') if isinstance(x, str) else x))
+
+        print("Replacing ZWJ token")
+        for column in result_col_names:
+            results_df[column] = PrePostProcessingUtils.clean_zwj(results_df[column])
+
+        PrePostProcessingUtils.save_dataframe(results_df, output_dir)
+
+        return results_df
+
+    def correctFromFile(self, src: str, target: str = None, output_dir: str = "outputs", max_length: int = 128,
+                        batch_size: int = 8, shuffle: bool = False):
+        """
+        Corrects the text from a file. The file should be in the format of
+        a CSV or TXT file. If the file is a CSV, it should have two columns: 'text' and 'expected'.
+        If the file is a TXT, it should have one line per sentence. If target is provided,
+        it will be used for evaluation. The output will be saved in the specified output directory.
+        """
+
+        if src.endswith('.csv'):
+            df = pd.read_csv(src)
+            input_col, target_col = ConfigUtils.get_dataset_columns()
+            src = df[input_col].tolist()
+            if target_col in df.columns:
+                target = df[target_col].tolist()
+
+        elif src.endswith('.txt'):
+            with open(src, 'r', encoding='utf-8') as file:
+                src = file.readlines()
+            if target is not None:
+                with open(target, 'r', encoding='utf-8') as file:
+                    target = file.readlines()
+        else:
+            raise LMSpellException("Unsupported file format. Only .csv and .txt are supported.") from None
+
+        return self.correct(src, target, output_dir, max_length, batch_size, shuffle)
+
+    def load_model(self):
+        pass
+
+    def save_model(self):
+        pass
+
+    def get_model_name(self):
+        pass
+
+    def load_tokenizer(self):
+        pass
+
+    def save_tokenizer(self):
+        pass