erictransformer 0.0.1__py3-none-any.whl

erictransformer/eric_tasks/eric_text_to_text.py

@@ -0,0 +1,283 @@
+import textwrap
+import threading
+from dataclasses import dataclass
+from typing import Iterator, List, Optional, Tuple, Union
+
+from datasets import Dataset
+from transformers import (
+    AutoModelForSeq2SeqLM,
+    DataCollatorForSeq2Seq,
+    GenerationConfig,
+    PretrainedConfig,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    PreTrainedTokenizerBase,
+    PreTrainedTokenizerFast,
+    TextIteratorStreamer,
+)
+
+from erictransformer.args import EricTrainArgs, EricEvalArgs
+from erictransformer.eval_models import EvalModel
+from erictransformer.exceptions import EricTokenizationError
+from erictransformer.eric_tasks.args import (
+    TTCallArgs,
+    TTTokArgs,
+)
+from erictransformer.eric_tasks.misc import generate_tt_kwargs, get_pad_eos
+from erictransformer.eric_tasks.results import TTResult
+from erictransformer.eric_tasks.tok.tok_functions import get_max_in_len
+from erictransformer.eric_transformer import EricTransformer, EricTransformerArgs
+from erictransformer.loops import EvalResult
+from erictransformer.utils import get_model_components
+from erictransformer.validator import TTValidator
+
+
+@dataclass(kw_only=True)
+class TTStreamResult:
+    text: str
+
+
+class EricTextToText(EricTransformer):
+    def __init__(
+        self,
+        model_name: Union[str, PreTrainedModel, None] = "google-t5/t5-base",
+        *,
+        trust_remote_code: bool = False,
+        tokenizer: Union[str, PreTrainedTokenizer, PreTrainedTokenizerFast] = None,
+    ):
+        model_class = AutoModelForSeq2SeqLM
+        eric_args = EricTransformerArgs(
+            model_name=model_name,
+            model_class=model_class,
+            trust_remote_code=trust_remote_code,
+            tokenizer=tokenizer
+        )
+        super().__init__(eric_args)
+
+        self.task_validator = TTValidator(
+            model_name=model_name,
+            trust_remote_code=trust_remote_code,
+            tokenizer=tokenizer,
+            logger=self.logger,
+        )
+
+        if self.model is not None:
+            self.pad_token_id, self.eos_token_id = get_pad_eos(
+                self.tokenizer, self.model
+            )
+
+            self._prep_model()
+            self._data_collator = DataCollatorForSeq2Seq(
+                self.tokenizer, model=self.model
+            )
+
+    def _get_call_thread_streamer(self, text: str, args: TTCallArgs = TTCallArgs()):
+        tokens = self.tokenizer(text, return_tensors="pt", truncation=True).to(
+            self.device
+        )
+        input_ids = tokens["input_ids"]
+        attention_mask = tokens["attention_mask"]
+        if input_ids.ndim == 1:
+            input_ids = input_ids.unsqueeze(0)
+
+        gen_streamer = TextIteratorStreamer(
+            self.tokenizer, skip_prompt=True, skip_special_tokens=True
+        )
+
+        gen_kwargs = generate_tt_kwargs(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            streamer=gen_streamer,
+            args=args,
+            eos_token_id=self.eos_token_id,
+            pad_token_id=self.pad_token_id
+        )
+
+        gen_thread = threading.Thread(target=self.model.generate, kwargs=gen_kwargs)
+
+        return gen_thread, gen_streamer
+
+    def __call__(
+        self,
+        text: str,
+        args: TTCallArgs = TTCallArgs(),
+    ) -> TTResult:
+        self._get_model_ready_inference()
+        self.task_validator.validate_call(text, args)
+        gen_thread, gen_streamer = self._get_call_thread_streamer(text, args)
+        gen_thread.start()
+        out_text = []
+        try:
+            for text in gen_streamer:
+                out_text.append(text)
+        finally:
+            gen_thread.join()
+            pass
+
+        final_text = "".join(out_text)
+        return TTResult(text=final_text)
+
+    def stream(
+        self, text: str, args: TTCallArgs = TTCallArgs()
+    ) -> Iterator[TTStreamResult]:
+        self._get_model_ready_inference()
+        self.task_validator.validate_call(text, args)
+
+        gen_thread, gen_streamer = self._get_call_thread_streamer(text, args)
+        gen_thread.start()
+        try:
+            for text in gen_streamer:
+                yield TTStreamResult(text=text)
+        finally:
+            gen_thread.join()
+
+    def _tok_function(
+        self,
+        raw_dataset,
+        args: TTTokArgs = TTTokArgs(),
+        file_type: str = "",
+        procs: Optional[int] = None,
+    ) -> Dataset:
+        final_max_in_len = get_max_in_len(
+            args.max_in_len, self.tokenizer
+        )
+        final_max_out_len = get_max_in_len(
+            args.max_out_len, self.tokenizer
+        )
+
+        def __preprocess_function(examples):
+            try:
+                model_inputs = self.tokenizer(
+                    examples["input"],
+                    max_length=final_max_in_len,
+                    truncation=True,
+                    padding="max_length",
+                )
+
+                labels = self.tokenizer(
+                    examples["target"],
+                    max_length=final_max_out_len,
+                    truncation=True,
+                    padding="max_length",
+                )
+
+                model_inputs["labels"] = [
+                    [
+                        (tok if tok != self.tokenizer.pad_token_id else -100)
+                        for tok in seq
+                    ]
+                    for seq in labels["input_ids"]
+                ]
+                return model_inputs
+            except Exception as e:
+                raise EricTokenizationError(
+                    f"Tokenization failed during preprocessing: {e}"
+                )
+
+        try:
+            tok_dataset = raw_dataset.map(
+                __preprocess_function,
+                batched=True,
+                remove_columns=["input", "target"],
+                batch_size=args.bs,
+                desc="Tokenizing...",
+                num_proc=procs,
+            )
+            tok_dataset.set_format(
+                type="torch", columns=["input_ids", "attention_mask", "labels"]
+            )
+            return tok_dataset
+        except Exception as e:
+            raise EricTokenizationError(
+                f"Failed to apply preprocessing function over dataset: {e}"
+            )
+
+    def train(
+        self,
+        train_path: str = "",
+        args: EricTrainArgs = EricTrainArgs(),
+        eval_path: str = "",
+        resume_path: str = "",
+    ):
+        return super().train(train_path, args, eval_path, resume_path=resume_path)
+
+    def eval(self, eval_path: str = "", args=EricEvalArgs()) -> EvalResult:
+        return super().eval(eval_path=eval_path, args=args)
+
+    def tok(
+        self,
+        path: str,
+        out_dir: str,
+        args: TTTokArgs = TTTokArgs()
+    ):
+        return super().tok(path=path, out_dir=out_dir, args=args)
+
+    def _load_model_components(
+        self,
+    ) -> Tuple[PretrainedConfig, PreTrainedTokenizerBase, PreTrainedModel]:
+        return get_model_components(
+            model_name_path=self.eric_args.model_name,
+            trust_remote_code=self.eric_args.trust_remote_code,
+            model_class=self.eric_args.model_class,
+            tokenizer_path=self.eric_args.tokenizer,
+            precision=self.precision_type,
+        )
+
+    def _format_tokenized_example(self, example: dict) -> dict:
+        return {
+            "input_ids": example["input_ids"],
+            "attention_mask": example["attention_mask"],
+            "labels": example["labels"],
+        }
+
+    def _get_default_eval_models(self) -> List[EvalModel]:
+        return []
+
+    def _prep_model(self):
+        generation_config = GenerationConfig.from_model_config(self.model.config)
+        args = TTCallArgs()
+        generation_config.num_beams = 1
+        generation_config.early_stopping = False
+        generation_config.do_sample = True
+        generation_config.min_len = args.min_len
+        generation_config.max_length = args.max_len
+        generation_config.temp = args.temp
+        generation_config.top_p = args.top_p
+        self.model.generation_config = generation_config
+
+    def _get_readme(self, repo_id: str) -> str:
+        readme_text = textwrap.dedent(f"""\
+        ---
+        tags:
+        - erictransformer
+        - eric-text-to-text
+        ---
+        # {repo_id}
+
+        ## Installation
+
+        ```
+        pip install erictransformer
+        ```
+
+        ## Usage 
+
+        ```python
+        from erictransformer import EricTextToText, TTCallArgs
+
+        eric_tt = EricTextToText(model_name="{repo_id}")
+
+        text = 'Hello world'
+        
+        result = eric_tt(text)
+        print(result.text)
+        
+        # Stream
+        for chunk in eric_tt.stream(text):
+            print(chunk.text, end="")
+        ```
+        
+        See Eric Transformer's [GitHub](https://github.com/ericfillion/erictransformer) for more information. 
+        """)
+
+        return readme_text

erictransformer/eric_tasks/inference_engine/__init__.py

@@ -0,0 +1,3 @@
+from erictransformer.eric_tasks.inference_engine.text_classification import (
+    tc_inference,
+)

erictransformer/eric_tasks/inference_engine/text_classification.py

@@ -0,0 +1,28 @@
+from typing import List, Tuple
+
+import torch
+
+
+@torch.inference_mode()
+def tc_inference(
+    tokens,
+    model,
+    id2label,
+) -> List[List[Tuple[str, float]]]:
+    input_ids: torch.Tensor = tokens["input_ids"]
+    attention_mask: torch.Tensor = tokens["attention_mask"]
+    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+    logits = outputs.logits
+    probs = torch.softmax(logits, dim=-1)
+
+    scores, indices = probs.sort(dim=-1, descending=True)
+
+    output: List[List[Tuple[str, float]]] = [
+        [
+            (id2label[int(i)], float(s))
+            for i, s in zip(idx_row.tolist(), sc_row.tolist())
+        ]
+        for idx_row, sc_row in zip(indices, scores)
+    ]
+
+    return output

erictransformer/eric_tasks/misc/__init__.py

@@ -0,0 +1,11 @@
+from erictransformer.eric_tasks.misc.call_utils import (
+    format_messages,
+    generate_gen_kwargs,
+    generate_tt_kwargs,
+)
+from erictransformer.eric_tasks.misc.get_pad_eos import get_pad_eos
+from erictransformer.eric_tasks.misc.rag import (
+    create_search_prompt_chat,
+    formate_rag_content,
+    formate_rag_message,
+)

erictransformer/eric_tasks/misc/call_utils.py

@@ -0,0 +1,69 @@
+from typing import Union
+
+from transformers import TextIteratorStreamer
+
+from erictransformer.exceptions import EricInferenceError
+from erictransformer.eric_tasks.args import CHATCallArgs, GENCallArgs, TTCallArgs
+
+
+def format_messages(text):
+    if isinstance(text, str):
+        messages = [{"role": "user", "content": text}]
+    elif isinstance(text, list) and all(isinstance(el, dict) for el in text):
+        messages = text
+    else:
+        raise EricInferenceError("Wrong input format")
+
+    return messages
+
+
+def generate_gen_kwargs(
+    input_ids,
+    attention_mask,
+    streamer: TextIteratorStreamer,
+    args: Union[CHATCallArgs, GENCallArgs],
+    eos_token_id: int,
+    pad_token_id: int,
+) -> dict:
+    max_len = args.max_len
+    if args.min_len > args.max_len:
+        max_len = args.min_len
+
+    gen_kwargs = dict(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        streamer=streamer,
+        max_new_tokens=max_len,
+        temp=args.temp,
+        top_p=args.top_p,
+        eos_token_id=eos_token_id,
+        pad_token_id=pad_token_id,
+    )
+
+    return gen_kwargs
+
+
+def generate_tt_kwargs(
+    input_ids,
+    attention_mask,
+    streamer: TextIteratorStreamer,
+    args: Union[CHATCallArgs, TTCallArgs],
+    eos_token_id: int,
+    pad_token_id: int
+) -> dict:
+    max_len = args.max_len
+    if args.min_len > args.max_len:
+        max_len = args.min_len
+
+    gen_kwargs = dict(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        streamer=streamer,
+        max_new_tokens=max_len,
+        temp=args.temp,
+        top_p=args.top_p,
+        eos_token_id=eos_token_id,
+        pad_token_id=pad_token_id
+    )
+
+    return gen_kwargs

erictransformer/eric_tasks/misc/get_pad_eos.py

@@ -0,0 +1,24 @@
+from transformers import PreTrainedModel, PreTrainedTokenizer
+
+from erictransformer.exceptions import EricInferenceError
+
+
+def get_pad_eos(tokenizer: PreTrainedTokenizer, model: PreTrainedModel):
+    if tokenizer.pad_token_id is not None:
+        pad_id = tokenizer.pad_token_id
+    elif tokenizer.eos_token_id is not None:
+        pad_id = tokenizer.eos_token_id
+    else:
+        raise EricInferenceError(
+            "Tokenizer doesn't have a pad_token_id or eos_token_id token"
+        )
+
+    if model.config.eos_token_id is not None:
+        eos_id = model.config.eos_token_id
+    elif tokenizer.eos_token_id is not None:
+        eos_id = tokenizer.eos_token_id
+    else:
+        raise EricInferenceError(
+            "The model and the tokenizer don't define an eos_token_id"
+        )
+    return pad_id, eos_id

erictransformer/eric_tasks/misc/rag.py

@@ -0,0 +1,17 @@
+from ericsearch import RankerResult
+
+
+def formate_rag_content(text: str, data_result: RankerResult):
+    return f"Based on the search query: '{text}'. The following data may be relevant: ' {data_result.text} '"
+
+
+def formate_rag_message(rag_content: str):
+    return {"role": "user", "content": rag_content}
+
+
+def create_search_prompt_chat(text: str) -> str:
+    if type(text) == str:
+        search_query = text
+    else:
+        search_query = text[-1]["content"]
+    return search_query

erictransformer/eric_tasks/results/__init__.py

@@ -0,0 +1,6 @@
+from erictransformer.eric_tasks.results.call_results import (
+    CHATResult,
+    GENResult,
+    TCResult,
+    TTResult
+)

erictransformer/eric_tasks/results/call_results.py

@@ -0,0 +1,30 @@
+from dataclasses import dataclass
+from typing import List
+
+from erictransformer.args import CallResult
+
+
+@dataclass(kw_only=True)
+class GENResult(CallResult):
+    text: str
+
+
+@dataclass(kw_only=True)
+class CHATResult(CallResult):
+    text: str
+
+
+@dataclass(kw_only=True)
+class StreamCHATResult(CallResult):
+    text: str
+    mode: str
+
+
+@dataclass(kw_only=True)
+class TCResult(CallResult):
+    labels: List[str]
+    scores: List[float]
+
+@dataclass(kw_only=True)
+class TTResult(CallResult):
+    text: str

erictransformer/eric_tasks/tok/tok_functions.py

@@ -0,0 +1,118 @@
+from typing import Union
+
+from datasets import Dataset
+from transformers import PreTrainedTokenizer
+
+from erictransformer.exceptions import EricTokenizationError
+
+
+def get_max_in_len(
+    max_len: Union[int, None], tokenizer: PreTrainedTokenizer
+) -> int:
+    if max_len >=1:
+        return max_len
+    else:
+        if tokenizer.model_max_length > 10_000_000:
+            return 512
+        return tokenizer.model_max_length
+
+
+def tokenize_gen(
+    tokenizer: PreTrainedTokenizer,
+    dataset: Dataset,
+    max_len: int,
+    bs: int,
+    procs: int = 1,
+) -> Dataset:
+    def tokenize_fn(batch):
+        try:
+            tokens = tokenizer(
+                batch["text"],
+                padding="max_length",
+                truncation=True,
+                max_length=max_len,
+            )
+            labels = []
+            for input_ids in tokens["input_ids"]:
+                masked = [
+                    token if token != tokenizer.pad_token_id else -100
+                    for token in input_ids
+                ]
+                labels.append(masked)
+            tokens["labels"] = labels
+            return tokens
+        except Exception as e:
+            raise EricTokenizationError(
+                f"Tokenization failed during batch processing. Error: {e}"
+            )
+
+    try:
+        tokenized = dataset.map(
+            tokenize_fn,
+            batched=True,
+            remove_columns=["text"],
+            batch_size=bs,
+            desc="Tokenizing...",
+            num_proc=procs,
+        )
+        cols = ["input_ids", "attention_mask"]
+        cols.append("labels")
+
+        tokenized.set_format("torch", columns=cols)
+        return tokenized
+    except Exception as e:
+        raise EricTokenizationError(
+            f"Failed during dataset mapping or formatting. Error: {e}"
+        )
+
+
+def tokenize_chat_template(
+    tokenizer: PreTrainedTokenizer,
+    dataset: Dataset,
+    max_len: int,
+    bs: int,
+    procs: int = 1,
+) -> Dataset:
+    def tokenize_fn(example):
+        try:
+            inputs = [msg for msg in example["messages"]]
+            tokens = tokenizer.apply_chat_template(
+                inputs,
+                tokenize=True,
+                add_generation_prompt=False,
+                max_length=max_len,
+                padding="max_length",
+                return_dict=True,
+                truncation=True,
+            )
+
+            labels = []
+            for input_ids in tokens["input_ids"]:
+                masked = [
+                    token if token != tokenizer.pad_token_id else -100
+                    for token in input_ids
+                ]
+                labels.append(masked)
+            tokens["labels"] = labels
+
+            return tokens
+        except Exception as e:
+            raise EricTokenizationError(
+                f"Tokenization failed during chat template application. Error: {e}"
+            )
+
+    try:
+        tokenized = dataset.map(
+            tokenize_fn,
+            batched=True,
+            remove_columns=["messages"],
+            batch_size=bs,
+            desc="Tokenizing...",
+            num_proc=procs,
+        )
+        tokenized.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
+        return tokenized
+    except Exception as e:
+        raise EricTokenizationError(
+            f"Failed during dataset mapping or formatting. Error: {e}"
+        )

erictransformer/eric_tracker/__init__.py

@@ -0,0 +1 @@
+from erictransformer.eric_tracker.eric_tracker import EricTracker