Delete srl_pipeline.py

srl_pipeline.py

@@ -1,242 +0,0 @@
-import logging
-from typing import Any, Dict, List, Tuple
-
-import spacy
-import torch
-from transformers import Pipeline
-
-from decoder import Decoder
-
-logger = logging.getLogger(__name__)
-
-
-class SrlPipeline(Pipeline):
-    """
-    A pipeline for Semantic Role Labeling (SRL) using transformers and spaCy models.
-
-    This pipeline tokenizes input sentences, finds verbs using POS tagging, and postprocesses
-    the model outputs using Viterbi decoding to provide human-readable results.
-
-    Attributes:
-        model ``str``: The name or identifier of the underlying transformer model.
-        tokenizer ``str``: The name or identifier of the tokenizer associated with the model.
-        framework ``str``: The framework used for the pipeline (e.g., PyTorch, TensorFlow).
-        task ``str``: The specific task of the pipeline.
-        verb_predictor: An instance of spaCy model used for predicting verbs in the input sentences.
-    Usage:
-        # Register the SrlPipeline in the pipeline registry
-        PIPELINE_REGISTRY.register_pipeline(
-            "srl",
-            pipeline_class=SrlPipeline,
-            model=SRLModel,  # Assuming SRLModel is the model class used
-            default={"lang": "en"},
-            type="text",
-        )
-
-        # Load the model and tokenizer
-        model = AutoModel.from_pretrained("liaad/srl-en_roberta-large_hf", trust_remote_code=True)
-        tokenizer = AutoTokenizer.from_pretrained("liaad/srl-en_roberta-large_hf", trust_remote_code=True)
-
-        # Load the SRL pipeline
-        srl_pipeline = pipeline(
-            "srl",
-            model=model,
-            tokenizer=tokenizer,
-            framework="PyTorch",  # Replace with actual framework used
-            task="semantic_role_labeling",  # Replace with actual task name
-            lang="en"  # Language specification
-        )
-
-        # Example text input
-        text = ["The cat jumps over the fence.", "She quickly eats the delicious cake."]
-
-        # Perform semantic role labeling
-        results = srl_pipeline(text)
-    """
-
-    def __init__(self, model: str, tokenizer: str, framework: str, task: str, **kwargs):
-        """
-        Initializes the Semantic Role Labeling pipeline.
-
-        Parameters:
-        - model ``str``: The model name or identifier.
-        - tokenizer ``str``: The tokenizer name or identifier.
-        - framework ``str``: The framework used.
-        - task ``str``: The specific task of the pipeline.
-        - **kwargs: Additional keyword arguments.
-                - lang ``str``, optional: Language specification ('en' for English or 'pt' for Portuguese, which is default).
-        """
-        super().__init__(model, tokenizer=tokenizer)
-        if "lang" in kwargs and kwargs["lang"] == "en":
-            logger.info("Loading English verb predictor model...")
-            self.verb_predictor = spacy.load("en_core_web_trf")
-        else:
-            logger.info("Loading Portuguese verb predictor model...")
-            self.verb_predictor = spacy.load("pt_core_news_lg")
-        logger.info("Got verb prediction model\n")
-
-    def _sanitize_parameters(
-        self, **kwargs: Dict[str, Any]
-    ) -> Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]:
-        """
-        Sanitizes and organizes additional parameters.
-
-        Parameters:
-        - **kwargs: Additional keyword arguments.
-
-        Returns:
-        - ``Tuple[Dict[str, Any], Dict[str, Any], Dict[str, Any]]``: Three dictionaries of sanitized parameters for preprocess, _forward, and postprocess.
-        """
-        return {}, {}, {}
-
-    def preprocess(self, sentence: str) -> List[Dict[str, Any]]:
-        """
-        Preprocesses a sentence for semantic role labeling.
-
-        Parameters:
-        - sentence ``str``: The input sentence to be processed.
-
-        Returns:
-        - ``List[Dict[str, Any]]``: A list of dictionaries containing model inputs for each verb in the sentence.
-        """
-        # Extract sentence verbs
-        doc = self.verb_predictor(sentence)
-
-        verbs = {token.text for token in doc if token.pos_ == "VERB"}
-        # If the sentence only contains auxiliary verbs, consider those as the
-        # main verbs
-        if not verbs:
-            verbs = {token.text for token in doc if token.pos_ == "AUX"}
-
-        # Tokenize sentence
-        tokens = self.tokenizer.encode_plus(
-            sentence,
-            truncation=True,
-            return_token_type_ids=False,
-            return_offsets_mapping=True,
-        )
-        tokens_lst = tokens.tokens()
-        offsets = tokens["offset_mapping"]
-
-        input_ids = torch.tensor([tokens["input_ids"]], dtype=torch.long)
-        attention_mask = torch.tensor([tokens["attention_mask"]], dtype=torch.long)
-
-        model_input = {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask,
-            "token_type_ids": [],
-            "tokens": tokens_lst,
-            "verb": "",
-        }
-
-        model_inputs = [
-            {**model_input} for _ in verbs
-        ]  # Create a new dictionary for each verb
-
-        for i, verb in enumerate(verbs):
-            model_inputs[i]["verb"] = verb
-            token_type_ids = model_inputs[i]["token_type_ids"]
-            token_type_ids.append([])
-            curr_word_offsets: tuple[int, int] = None
-
-            for j in range(len(tokens_lst)):
-                curr_offsets = offsets[j]
-                curr_slice = sentence[curr_offsets[0] : curr_offsets[1]]
-                if not curr_slice:
-                    token_type_ids[-1].append(0)
-                # Check if new token still belongs to same word
-                elif (
-                    curr_word_offsets
-                    and curr_offsets[0] >= curr_word_offsets[0]
-                    and curr_offsets[1] <= curr_word_offsets[1]
-                ):
-                    # Extend previous token type
-                    token_type_ids[-1].append(token_type_ids[-1][-1])
-                else:
-                    curr_word_offsets = self._find_word(sentence, start=curr_offsets[0])
-                    curr_word = sentence[curr_word_offsets[0] : curr_word_offsets[1]]
-
-                    token_type_ids[-1].append(
-                        int(curr_word != "" and curr_word == verb)
-                    )
-
-            model_inputs[i]["token_type_ids"] = torch.tensor(
-                token_type_ids, dtype=torch.long
-            )
-
-        return model_inputs
-
-    def _forward(self, model_inputs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-        """
-        Internal method to forward model inputs for prediction.
-
-        Parameters:
-        - model_inputs ``List[Dict[str, Any]]``: List of dictionaries containing model inputs.
-
-        Returns:
-        - ``List[Dict[str, Any]]``: List of dictionaries containing model outputs.
-        """
-        outputs = []
-        for model_input in model_inputs:
-            output = self.model(
-                input_ids=model_input["input_ids"],
-                attention_mask=model_input["attention_mask"],
-                token_type_ids=model_input["token_type_ids"],
-            )
-            output["verb"] = model_input["verb"]
-            output["tokens"] = model_input["tokens"]
-            outputs.append(output)
-        return outputs
-
-    def postprocess(self, model_outputs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-        """
-        Postprocesses model outputs to human-readable format.
-
-        Parameters:
-        - model_outputs ``List[Dict[str, Any]]``: List of dictionaries containing model outputs.
-
-        Returns:
-        - ``List[Dict[str, Any]]``: List of dictionaries containing processed results.
-                Each dictionary entry represents a verb with its associated labels and token-label pairs.
-                Example format: {verb: (labels, List[(token, label)])}
-        """
-        result = []
-        id2label = {int(k): str(v) for k, v in self.model.config.id2label.items()}
-        evaluator = Decoder(id2label)
-
-        for model_output in model_outputs:
-            class_probabilities = model_output["class_probabilities"]
-            attention_mask = model_output["attention_mask"]
-            output_dict = evaluator.make_output_human_readable(
-                class_probabilities, attention_mask
-            )
-            # Here we always fetch the first list because in a pipeline every
-            # sentence is processed one at a time
-            wordpiece_label_ids = output_dict["wordpiece_label_ids"][0]
-            labels = list(map(lambda idx: id2label[idx], wordpiece_label_ids))
-            result.append(
-                {
-                    model_output["verb"]: (
-                        labels,
-                        list(zip(model_output["tokens"], labels)),
-                    )
-                }
-            )
-        return result
-
-    def _find_word(self, s: str, start: int = 0) -> Tuple[int, int]:
-        """
-        Helper method to find the boundaries of a word in a string.
-        Assumes a non alphanumeric char represents the end of a word.
-
-        Parameters:
-        - s ``str``: The input string.
-        - start ``int``, optional: Starting index to start looking for the word. Defaults to 0.
-
-        Returns:
-        - ``Tuple[int, int]``: Start and end indices of the word.
-        """
-        for i, char in enumerate(s[start:], start):
-            if not char.isalpha():
-                return start, i
-        return start, len(s)