From 039d71830609d77024718dfd4dbd05d51d00e95f Mon Sep 17 00:00:00 2001
From: RayanMarmar <106345540+RayanMarmar@users.noreply.github.com>
Date: Tue, 19 Mar 2024 18:54:02 +0100
Subject: [PATCH 1/3] Hotfix transformers class (#79)

* fix bug transformers class name

* add default class

* fixed unit tests
---
 sdk/downloader.py            | 9 ++++-----
 sdk/tests/test_downloader.py | 6 +++---
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/sdk/downloader.py b/sdk/downloader.py
index c4fd88d..04077f7 100644
--- a/sdk/downloader.py
+++ b/sdk/downloader.py
@@ -254,12 +254,11 @@ def set_transformers_class_names(model: Model) -> None:
         # Get the configuration
         config = transformers.AutoConfig.from_pretrained(model.name)
 
-        # Map model class from model type
-        model_mapping = transformers.AutoModel._model_mapping._model_mapping
-
         # Set model class name if not already set
-        model.class_name = model.class_name or model_mapping.get(
-            config.model_type)
+        model.class_name = model.class_name or config.architectures[0] \
+            if config.architectures and config.architectures[0] else (
+            model_config_default_class_for_module[TRANSFORMERS]
+        )
 
         # Set tokenizer class name if not already set
         # and config.tokenizer_class exists
diff --git a/sdk/tests/test_downloader.py b/sdk/tests/test_downloader.py
index 2a5a1eb..a72e0c3 100644
--- a/sdk/tests/test_downloader.py
+++ b/sdk/tests/test_downloader.py
@@ -893,7 +893,7 @@ def test_set_diffusers_class_names_with_configured_model(self,
         self.assertEqual(model.class_name, 'TestModel')
 
     @patch('transformers.AutoConfig.from_pretrained',
-           return_value=MagicMock(model_type='t5',
+           return_value=MagicMock(architectures=['T5Model'],
                                   tokenizer_class='TokenizerClass'))
     def test_set_transformers_class_names(self, mock_load_config):
         # Init
@@ -909,7 +909,7 @@ def test_set_transformers_class_names(self, mock_load_config):
         self.assertEqual(model.tokenizer.class_name, 'TokenizerClass')
 
     @patch('transformers.AutoConfig.from_pretrained',
-           return_value=MagicMock(model_type='t5',
+           return_value=MagicMock(architectures=['T5Model'],
                                   tokenizer_class=None))
     def test_set_transformers_class_names_with_default_tokenizer(
             self, mock_load_config
@@ -927,7 +927,7 @@ def test_set_transformers_class_names_with_default_tokenizer(
         self.assertEqual(model.tokenizer.class_name, 'AutoTokenizer')
 
     @patch('transformers.AutoConfig.from_pretrained',
-           return_value=MagicMock(model_type='t5',
+           return_value=MagicMock(architectures=['T5Model'],
                                   tokenizer_class='TokenizerClass'))
     def test_set_transformers_class_names_with_configured_model(
             self, mock_load_config

From c1a9580e0c0cccba92f579a3ba529f54b1d2450b Mon Sep 17 00:00:00 2001
From: Quentin <63395737+SkelNeXus@users.noreply.github.com>
Date: Wed, 20 Mar 2024 16:03:57 +0100
Subject: [PATCH 2/3] 78 sdk correction on the text conversation model (#81)

* Modification of the model text conversation, this code must be test

* Correction of the naming

* Creation of the model transformers class, modification of all model class and demo, need test

* Fix on modules and args

* fix issues on the text conv

* Suppression of options class and some corrections

* fix model.py

* fix main.py

* fix option.py

* fix demo_text_conv
---
 main.py                                       |  12 +-
 sdk/demo/__init__.py                          |   2 +-
 sdk/demo/demo_text_conv.py                    |  88 +-----
 ...demo_txt_to_img.py => demo_text_to_img.py} |  30 +-
 sdk/models/__init__.py                        |   1 +
 sdk/models/model.py                           |   7 +-
 sdk/models/model_text_conversation.py         | 282 +++---------------
 sdk/models/model_text_to_image.py             |  70 +----
 sdk/models/model_transformers.py              | 146 +++++++++
 sdk/models/models_management.py               |  68 +----
 sdk/options/__init__.py                       |   4 +-
 sdk/options/options.py                        |  16 -
 sdk/options/options_text_conversation.py      | 114 -------
 sdk/options/options_text_to_image.py          | 264 ----------------
 sdk/options/options_tokenizer.py              |  33 --
 sdk/tokenizers/tokenizer.py                   |  62 ++--
 16 files changed, 283 insertions(+), 916 deletions(-)
 rename sdk/demo/{demo_txt_to_img.py => demo_text_to_img.py} (53%)
 create mode 100644 sdk/models/model_transformers.py
 delete mode 100644 sdk/options/options_text_conversation.py
 delete mode 100644 sdk/options/options_text_to_image.py
 delete mode 100644 sdk/options/options_tokenizer.py

diff --git a/main.py b/main.py
index 808caee..7f8850a 100644
--- a/main.py
+++ b/main.py
@@ -1,21 +1,21 @@
 import argparse
-from sdk.demo import DemoTextConv, DemoTxtToImg
+from sdk.demo import DemoTextConv, DemoTextToImg
 
 
 if __name__ == '__main__':
 
     parser = argparse.ArgumentParser(description='Choose a Model type :'
-                                                 ' TextConv or TxtToImg')
+                                                 ' TextConv or TextToImg')
     subparser = parser.add_subparsers(dest='option')
 
     conv = subparser.add_parser('TextConv')
-    img = subparser.add_parser('TxtToImg')
+    img = subparser.add_parser('TextToImg')
     args = parser.parse_args()
 
     match args.option:
         case 'TextConv':
             DemoTextConv()
-        case ('TxtToImg'):
-            DemoTxtToImg()
+        case ('TextToImg'):
+            DemoTextToImg()
         case _:
-            DemoTxtToImg()
+            DemoTextToImg()
diff --git a/sdk/demo/__init__.py b/sdk/demo/__init__.py
index 2065acf..3b91259 100644
--- a/sdk/demo/__init__.py
+++ b/sdk/demo/__init__.py
@@ -1,3 +1,3 @@
 # flake8: noqa F401
 from .demo_text_conv import DemoTextConv
-from .demo_txt_to_img import DemoTxtToImg
+from .demo_text_to_img import DemoTextToImg
diff --git a/sdk/demo/demo_text_conv.py b/sdk/demo/demo_text_conv.py
index b0daaaf..65eb891 100644
--- a/sdk/demo/demo_text_conv.py
+++ b/sdk/demo/demo_text_conv.py
@@ -1,7 +1,6 @@
-from sdk.models import ModelsManagement, ModelsTextConversation
-from sdk.options import Devices, OptionsTextConversation
-from sdk.options.options_tokenizer import OptionsTokenizer
-from sdk.tokenizers.tokenizer import Tokenizer
+from sdk.models import ModelsTextConversation
+from sdk.options import Devices
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
 class DemoTextConv:
@@ -13,75 +12,20 @@ def __init__(self):
         """
         Initializes the DemoTextConv class with predefined options and models.
         """
-        # Define the model name and path
-        model_name = "Salesforce/codegen-350M-nl"
-        model_path = "Salesforce/codegen-350M-nl"
+        model_path = "microsoft/phi-2"
+        tokenizer_path = "microsoft/phi-2"
 
-        # Define options for text conversation
-        options = OptionsTextConversation(
-            prompt="Hello, what's 3 + 3 ?",
-            device=Devices.GPU,
-            model_name=model_name,
-            trust_remote_code=True,
-            max_length=100
-        )
+        model = ModelsTextConversation(model_name="model",
+                                       model_path=model_path,
+                                       tokenizer_path=tokenizer_path,
+                                       model_class=AutoModelForCausalLM,
+                                       tokenizer_class=AutoTokenizer,
+                                       device=Devices.GPU)
 
-        # Define tokenizer options
-        tokenizer_options = OptionsTokenizer(
-            device='cuda',
-            padding_side='left',
-            return_tensors="pt"
-        )
+        model.load_model()
+        model.create_new_conversation()
 
-        tokenizer = Tokenizer("Salesforce/codegen-350M-nl",
-                              "Salesforce/codegen-350M-nl",
-                              "Salesforce/codegen-350M-nl",
-                              tokenizer_options)
+        result = model.generate_prompt(
+            "I'm looking for a movie - what's your favourite one?")
 
-        # Initialize the model management
-        model_management = ModelsManagement()
-
-        # Create and load the text conversation model
-        model = ModelsTextConversation(model_name, model_path, options)
-        model.tokenizer = tokenizer
-        model_management.add_model(new_model=model, model_options=options)
-        model_management.load_model(model_name)
-
-        # Print the response to the initial prompt
-        print(model_management.generate_prompt(options.prompt))
-
-        # Generate a response to a custom prompt
-        print(model_management.generate_prompt(
-            prompt="What did I say before ?"))
-
-        # Create a new conversation with a new prompt
-        options = OptionsTextConversation(
-            prompt="Hello, what's 6 + 6 ?",
-            device=options.device,
-            model_name=model_name,
-            batch_size=1,
-            chat_id_to_use=1,
-            minimum_tokens=50,
-            create_new_conv=True
-        )
-
-        # Switch to the new conversation
-        model_management.set_model_options(model_name=model_name,
-                                           options=options)
-        print(model_management.generate_prompt(options.prompt))
-
-        print(model_management.generate_prompt("Where is Japan"))
-
-        # Switch back to the initial conversation
-        options.chat_id_to_use = 0
-        # Create a new tokenizer and use it
-        options.tokenizer_id_to_use = 1
-        model_management.set_model_options(model_name=model_name,
-                                           options=options)
-        tokenizer_options = OptionsTokenizer(
-            device='cuda',
-            padding_side='right',
-            return_tensors='pt'
-        )
-        model.tokenizer_options = tokenizer_options
-        print(model_management.generate_prompt("Bye "))
+        print(result.messages[-1]["content"])
diff --git a/sdk/demo/demo_txt_to_img.py b/sdk/demo/demo_text_to_img.py
similarity index 53%
rename from sdk/demo/demo_txt_to_img.py
rename to sdk/demo/demo_text_to_img.py
index 03d5d0a..bb253dd 100644
--- a/sdk/demo/demo_txt_to_img.py
+++ b/sdk/demo/demo_text_to_img.py
@@ -1,27 +1,29 @@
+import torch
 from sdk.models import ModelTextToImage, ModelsManagement
-from sdk.options import Devices, OptionsTextToImage
+from sdk.options import Devices
 
 
-class DemoTxtToImg:
+class DemoTextToImg:
 
     def __init__(self):
-        options = OptionsTextToImage(
-            prompt="Astronaut in a jungle, cold color palette, "
-                   "muted colors, detailed, 8k",
-            device=Devices.GPU,
-            image_width=512,
-            image_height=512,
-        )
-
         model_stabilityai_name = "stabilityai/sdxl-turbo"
         model_stabilityai_path = "stabilityai/sdxl-turbo"
         model_management = ModelsManagement()
         model_stabilityai = ModelTextToImage(model_stabilityai_name,
-                                             model_stabilityai_path)
+                                             model_stabilityai_path,
+                                             Devices.GPU,
+                                             torch_dtype=torch.float16,
+                                             use_safetensors=True,
+                                             add_watermarker=False,
+                                             variant="fp16")
 
-        model_management.add_model(new_model=model_stabilityai,
-                                   model_options=options)
+        model_management.add_model(new_model=model_stabilityai)
         model_management.load_model(model_stabilityai_name)
 
-        image = model_management.generate_prompt()
+        image = model_management.generate_prompt(
+            prompt="Astronaut in a jungle, cold color palette, "
+                   "muted colors, detailed, 8k",
+            image_width=512,
+            image_height=512
+        )
         image.show()
diff --git a/sdk/models/__init__.py b/sdk/models/__init__.py
index 6b44904..86152b0 100644
--- a/sdk/models/__init__.py
+++ b/sdk/models/__init__.py
@@ -1,5 +1,6 @@
 # flake8: noqa F401
 from sdk.models.model import Model
+from sdk.models.model_transformers import ModelTransformers
 from sdk.models.models_management import ModelsManagement
 from sdk.models.model_text_to_image import ModelTextToImage
 from sdk.models.model_text_conversation import ModelsTextConversation
diff --git a/sdk/models/model.py b/sdk/models/model.py
index 5e1d724..43a7e8f 100644
--- a/sdk/models/model.py
+++ b/sdk/models/model.py
@@ -1,6 +1,4 @@
 from abc import abstractmethod
-from sdk.options import Options
-from typing import Optional
 
 
 class Model:
@@ -20,7 +18,7 @@ def __init__(self, model_name, model_path: str):
         self.model_path = model_path
 
     @abstractmethod
-    def load_model(self, option: Options) -> bool:
+    def load_model(self) -> bool:
         raise NotImplementedError
 
     @abstractmethod
@@ -28,6 +26,5 @@ def unload_model(self) -> bool:
         raise NotImplementedError
 
     @abstractmethod
-    def generate_prompt(self, prompt: Optional[str],
-                        option: Options, **kwargs):
+    def generate_prompt(self, prompt: str, **kwargs):
         raise NotImplementedError
diff --git a/sdk/models/model_text_conversation.py b/sdk/models/model_text_conversation.py
index fc1da16..8a34c46 100644
--- a/sdk/models/model_text_conversation.py
+++ b/sdk/models/model_text_conversation.py
@@ -1,282 +1,90 @@
-import torch
-from typing import Optional, Dict
-from transformers import AutoModelForCausalLM
+import uuid
+from typing import Dict, Union, Any
+from transformers import (
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    Conversation
+)
 from sdk.tokenizers.tokenizer import Tokenizer
-from sdk.models import Model
-from sdk.options import Devices, OptionsTextConversation
+from sdk.models import ModelTransformers
+from sdk.options import Devices
 
 
-class ModelsTextConversation(Model):
+class ModelsTextConversation(ModelTransformers):
     """
     A class representing a text conversation model.
-
-    Attributes:
-        pipeline (AutoModelForCausalLM):
-            The pipeline for the text conversation.
-        tokenizer_object (Tokenizer):
-            The tokenizer object for the model.
-        tokenizer_options (OptionsTokenizer):
-            The options for the tokenizer.
-        tokenizer_dict (Dict[int, TokenizerObject]):
-            Dictionary to store tokenizers.
-        conversation_dict (Dict[int, Tuple[Conversation, list]]):
-            Dictionary to store conversations.
-        current_conversation_id (int):
-            ID of the current conversation.
-        current_tokenizer_id (int): ID of the current tokenizer.
-        conversation_ctr (int): Total number of conversations.
-        tokenizer_ctr (int): Total number of tokenizers.
-        loaded (bool): Flag indicating whether the model is loaded.
-        conversation_step (int): Step of the conversation.
-        conversation_history (list):
-            List to store chat history token IDs.
-        conversation_active (bool):
-         Flag indicating if a conversation is active.
     """
-    pipeline: AutoModelForCausalLM
     tokenizer: Tokenizer
-    options: OptionsTextConversation
-
-    tokenizer_dict: Dict[int, Tokenizer] = {}
-    conversation_dict: Dict[int, list] = {}
+    device: Union[str, Devices]
 
-    current_conversation_id: int = 0
-    current_tokenizer_id: int = 0
+    model_pipeline: PreTrainedModel
+    tokenizer_pipeline: PreTrainedTokenizer
+    conversation: Conversation
 
-    conversation_ctr: int = 0
-    tokenizer_ctr: int = 0
+    conversation_dict: Dict[uuid.UUID, Conversation] = {}
 
     loaded: bool
-    conversation_step: int = 0
-
-    conversation_active: bool = False
 
     def __init__(self, model_name: str, model_path: str,
-                 option: OptionsTextConversation):
+                 tokenizer_path,
+                 model_class: Any,
+                 tokenizer_class: Any,
+                 device: Union[str, Devices]
+                 ):
         """
         Initializes the ModelsTextToImage class
         :param model_name: The name of the model
         :param model_path: The path of the model
+        :param device: Which device the model must be on
         """
-        super().__init__(model_name, model_path)
-        self.options = option
-        self.loaded = False
-        self.create_pipeline()
-
-    def create_pipeline(self):
-        """
-        Creates the pipeline to load on the device
-        """
-        if self.loaded:
-            return
-        self.pipeline = AutoModelForCausalLM.from_pretrained(
-            self.model_path,
-            trust_remote_code=self.options.trust_remote_code,
-            pad_token_id=self.options.pad_token_id,
-            eos_token_id=self.options.eos_token_id,
-            max_length=self.options.max_length
-        )
-
-    def load_model(self, option: OptionsTextConversation) -> bool:
-        """
-        Load this model on the given device.
-
-        Args:
-            option (OptionsTextConversation): The options with the device.
-
-        Returns:
-            bool: True if the model is successfully loaded.
-        """
-
-        if self.loaded:
-            return True
-        if option.device == Devices.RESET:
-            return False
-        self.pipeline.to(option.device.value)
-        self.loaded = True
-        return True
-
-    def unload_model(self, option: OptionsTextConversation) -> bool:
-        """
-        Unloads the model
-        :return: True if the model is successfully unloaded
-        """
-        if not self.loaded:
-            return False
-        self.pipeline.to(device=(
-            option.device if isinstance(option.device, str) else (
-                option.device.value)))
-        torch.cuda.empty_cache()
-        torch.cuda.ipc_collect()
-        self.loaded = False
-        return True
+        super().__init__(model_name=model_name,
+                         model_path=model_path,
+                         tokenizer_path=tokenizer_path,
+                         task="conversational",
+                         model_class=model_class,
+                         tokenizer_class=tokenizer_class,
+                         device=device)
 
-    # Will be used to create new conversation
     def generate_prompt(
-            self, prompt: Optional[str],
-            option: OptionsTextConversation,
-            **kwargs):
+            self, prompt: str,
+            **kwargs) -> Union[Conversation, None]:
         """
         Generates the prompt with the given option.
 
         Args:
-            prompt (Optional[str]): The optional prompt.
-            option (OptionsTextConversation):
-                The options of text conversation model.
+            prompt: (str): The optional prompt.
 
         Returns:
             str: Generated prompt.
         """
-        prompt = prompt if prompt else option.prompt
-
-        if option.create_new_conv:
-            self.create_new_conversation(
-                option=option
-            )
-        if option.chat_id_to_use != self.current_conversation_id:
-            self.change_conversation(option.chat_id_to_use)
-
-        if option.tokenizer_id_to_use != self.current_tokenizer_id:
-            print("changing tokenizer")
-            self.set_tokenizer_to_use(option.tokenizer_id_to_use)
+        self.write_input(prompt)
+        return self.transformers_pipeline(self.conversation)
 
-        if not self.conversation_active:
-            print("creating conv")
-            self.create_new_conversation(
-                option=option
-            )
-
-        history = '\n'.join(
-            self.conversation_dict[self.current_conversation_id]
-        )
-        print("sending input")
-        str_to_send = self.tokenizer.prepare_input(prompt, history)
-        result = self.pipeline.generate(
-            str_to_send,
-            **kwargs
-        )
-        response = self.tokenizer.decode_model_output(result)
-
-        self.conversation_dict[self.current_conversation_id].append(prompt)
-        self.conversation_dict[self.current_conversation_id].append(response)
-        return response
-
-    """
-    Creates a new tokenizer.
-
-    Args:
-        tokenizer_options (TokenizerOptions): Options for the tokenizer.
-    """
-
-    def add_new_tokenizer(self,  # change to save_tokenizer
-                          tokenizer: Tokenizer) -> int:
-        # adding and ID to the tokenizer and saving it in a dict
-        self.current_tokenizer_id = self.tokenizer_ctr
-        self.tokenizer_dict[self.current_tokenizer_id] = tokenizer
-        self.tokenizer = tokenizer
-        self.tokenizer_ctr += 1
-        return self.current_tokenizer_id
-
-    def set_tokenizer_to_use(self, token_id: int) -> bool:
-        """
-        Set the tokenizer to use for generating text.
-
-        Args:
-            token_id (int): The ID of the tokenizer to use.
-
-        Returns:
-            int: The ID of the selected tokenizer,
-             or -1 if the provided ID is invalid.
-        """
-        if token_id in self.tokenizer_dict:
-            self.current_tokenizer_id = token_id
-            self.tokenizer = self.tokenizer_dict[token_id]
-            return True
-        else:
-            return False
-
-    def delete_tokenizer(self, token_id: int) -> bool:
-        """
-        Delete a tokenizer.
-
-        Args:
-            token_id (int):
-             The ID of the tokenizer to delete.
-
-        Returns:
-            int: 0 if the tokenizer was deleted successfully,
-             -1 if the provided ID is invalid.
-        """
-        if token_id in self.tokenizer_dict:
-            del self.tokenizer_dict[token_id]
-            return True
-        else:
-            return False
-
-    def delete_conversation(self, conversation_id: int) -> bool:
-        """
-        Delete a conversation.
-
-        Args:
-            conversation_id (int):
-                The ID of the conversation to delete.
-
-        Returns:
-            int: 0 if the conversation was deleted successfully,
-             -1 if the provided ID is invalid.
-        """
-        if conversation_id in self.conversation_dict:
-            del self.conversation_dict[conversation_id]
-            if len(self.conversation_dict) == 0:
-                self.conversation_active = False
-            return True
-        else:
-            return False
-
-    def send_new_input(self, prompt: str, history: list) -> str:
+    def write_input(self, prompt: str) -> None:
         """
         Send new input to the chatbot and generate a response.
 
         Args:
-            history: Chat history
             prompt (str): The input prompt for the chatbot.
 
         Returns:
             str: The generated response from the chatbot.
         """
-        input_ids = {"input_ids": self.tokenizer.prepare_input(
-            prompt, history=history)
-        }
-        return self.pipeline.generate(**input_ids)
+        # ToDo
+        schematic = ({"role": "user", "content": prompt})
+        self.conversation.add_message(schematic)
 
-    def create_new_conversation(self,
-                                option: OptionsTextConversation) -> int:
+    def create_new_conversation(self, **kwargs) -> None:
         """
         Create a new conversation.
-
-        Args:
-            option (OptionsTextConversation):
-                The options for the conversation.
-
-        Returns:
-            None
         """
-        option.create_new_conv = False
-        option.chat_id_to_use = self.conversation_ctr
-        self.conversation_active = True
-
-        # create new conversation and increment counter
-        self.current_conversation_id = self.conversation_ctr
-
-        # adding new conversation to dict, with cleared history
-        self.conversation_dict[
-            self.current_conversation_id
-        ] = []
-        # incrementing counter
-        self.conversation_ctr += 1
-        return self.current_conversation_id
+        conversation_uuid = uuid.uuid4()
+        conversation = Conversation(conversation_id=conversation_uuid,
+                                    **kwargs)
+        self.conversation_dict[conversation_uuid] = conversation
+        self.conversation = conversation
 
-    def change_conversation(self, conversation_id: int) -> bool:
+    def change_conversation(self, conversation_id: uuid.UUID) -> bool:
         """
         Change the active conversation.
 
@@ -289,7 +97,7 @@ def change_conversation(self, conversation_id: int) -> bool:
         """
         if conversation_id in self.conversation_dict:
             print("switching to conversation {}".format(conversation_id))
-            self.current_conversation_id = conversation_id
+            self.conversation = self.conversation_dict[conversation_id]
             return True
         else:
             return False
diff --git a/sdk/models/model_text_to_image.py b/sdk/models/model_text_to_image.py
index 950f42f..91e8ed6 100644
--- a/sdk/models/model_text_to_image.py
+++ b/sdk/models/model_text_to_image.py
@@ -1,7 +1,7 @@
 import torch
 from diffusers import DiffusionPipeline, StableDiffusionXLPipeline
-from sdk.options import OptionsTextToImage, Devices
-from typing import Optional
+from sdk.options import Devices
+from typing import Union
 from sdk.models import Model
 
 
@@ -11,18 +11,21 @@ class ModelTextToImage(Model):
     """
     pipeline: StableDiffusionXLPipeline
     loaded: bool
+    device: Union[str, Devices]
 
-    def __init__(self, model_name: str, model_path: str):
+    def __init__(self, model_name: str, model_path: str,
+                 device: Union[str, Devices], **kwargs):
         """
         Initializes the ModelsTextToImage class
         :param model_name: The name of the model
         :param model_path: The path of the model
         """
         super().__init__(model_name, model_path)
+        self.device = device
         self.loaded = False
-        self.create_pipeline()
+        self.create_pipeline(**kwargs)
 
-    def create_pipeline(self):
+    def create_pipeline(self, **kwargs):
         """
         Creates the pipeline to load on the device
         """
@@ -31,25 +34,21 @@ def create_pipeline(self):
 
         self.pipeline = DiffusionPipeline.from_pretrained(
             self.model_path,
-            torch_dtype=torch.float16,
-            use_safetensors=True,
-            add_watermarker=False,
-            variant="fp16",
+            **kwargs
         )
 
-    def load_model(self, option: OptionsTextToImage) -> bool:
+    def load_model(self) -> bool:
         """
         Load this model on the given device
-        :param option: The options with the device
         :return: True if the model is successfully loaded
         """
         if self.loaded:
             return True
-        if option.device == Devices.RESET:
+        if self.device == Devices.RESET:
             return False
         self.pipeline.to(device=(
-            option.device if isinstance(option.device, str) else (
-                option.device.value)))
+            self.device if isinstance(self.device, str) else (
+                self.device.value)))
         self.loaded = True
         return True
 
@@ -66,51 +65,14 @@ def unload_model(self) -> bool:
         self.loaded = False
         return True
 
-    def generate_prompt(self, prompt: Optional[str],
-                        options: OptionsTextToImage,
+    def generate_prompt(self, prompt: str,
                         **kwargs):
         """
         Generates the prompt with the given option
-        :param prompt: The optional prompt
-            (if the prompt is empty, the options.prompt will be used)
-        :param options: The options of text to image model
+        :param prompt: The prompt to generate
         :return: An object image resulting from the model
         """
         return self.pipeline(
-            prompt=prompt if prompt else options.prompt,
-            prompt_2=options.prompt_2,
-            width=options.image_width,
-            height=options.image_height,
-            num_inference_steps=options.num_inference_steps,
-            timesteps=options.timesteps,
-            denoising_end=options.denoising_end,
-            guidance_scale=options.guidance_scale,
-            negative_prompt=options.negative_prompt,
-            negative_prompt_2=options.negative_prompt_2,
-            num_images_per_prompt=options.num_images_per_prompt,
-            eta=options.eta,
-            generator=options.generator,
-            latents=options.latents,
-            prompt_embeds=options.prompt_embeds,
-            negative_prompt_embeds=options.negative_prompt_embeds,
-            pooled_prompt_embeds=options.pooled_prompt_embeds,
-            negative_pooled_prompt_embeds=(
-                                        options.negative_pooled_prompt_embeds),
-            ip_adapter_image=options.ip_adapter_image,
-            output_type=options.output_type,
-            return_dict=options.return_dict,
-            cross_attention_kwargs=options.cross_attention_kwargs,
-            guidance_rescale=options.guidance_rescale,
-            original_size=options.original_size,
-            crops_coords_top_left=options.crops_coords_top_left,
-            target_size=options.target_size,
-            negative_original_size=options.negative_original_size,
-            negative_crops_coords_top_left=(
-                options.negative_crops_coords_top_left),
-            negative_target_size=options.negative_target_size,
-            clip_skip=options.clip_skip,
-            callback_on_step_end=options.callback_on_step_end,
-            callback_on_step_end_tensor_inputs=(
-                options.callback_on_step_end_tensor_inputs),
+            prompt=prompt,
             **kwargs
         ).images[0]
diff --git a/sdk/models/model_transformers.py b/sdk/models/model_transformers.py
new file mode 100644
index 0000000..62f7d01
--- /dev/null
+++ b/sdk/models/model_transformers.py
@@ -0,0 +1,146 @@
+import torch
+from typing import Any, Union
+from transformers import (
+    pipeline,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+from sdk.models import Model
+from sdk.options import Devices
+
+
+class ModelTransformers(Model):
+    """
+    A class to use if the model is a transformers one
+    """
+    task: str
+    # ToDo: look for a best type
+    model_class: Any
+    tokenizer_class: Any
+    device: Union[str, Devices]
+
+    tokenizer_path: str
+
+    model_pipeline_args: dict[str, Any] = {}
+    tokenizer_pipeline_args: dict[str, Any] = {}
+
+    transformers_pipeline: pipeline
+    model_pipeline: PreTrainedModel
+    tokenizer_pipeline: PreTrainedTokenizer
+
+    loaded: bool
+
+    def __init__(self, model_name: str, model_path: str,
+                 tokenizer_path,
+                 task: str,
+                 model_class: Any,
+                 tokenizer_class: Any,
+                 device: Union[str, Devices]
+                 ):
+        """
+        Initializes the ModelsTextToImage class
+        :param model_name: The name of the model
+        :param model_path: The path of the model
+        :param device: Which device the model must be on
+        """
+        super().__init__(model_name, model_path)
+        self.tokenizer_path = tokenizer_path
+        self.task = task
+        self.device = device
+        self.loaded = False
+        self.model_class = model_class
+        self.tokeniser_class = tokenizer_class
+        self.create_pipeline()
+
+    def set_model_pipeline_args(self, **kwargs):
+        if kwargs:
+            self.model_pipeline_args = kwargs.copy()
+
+    def set_tokenizer_pipeline_args(self, **kwargs):
+        if kwargs:
+            self.tokenizer_pipeline_args = kwargs.copy()
+
+    def create_pipeline(self, **kwargs) -> None:
+        """
+        Creates the pipeline and load them on the device
+        """
+        if self.loaded:
+            return
+
+        self.model_pipeline = self.model_class.from_pretrained(
+            self.model_path,
+            **self.model_pipeline_args
+        )
+
+        self.tokenizer_pipeline = self.tokeniser_class.from_pretrained(
+            self.tokenizer_path,
+            **self.tokenizer_pipeline_args
+        )
+
+        self.transformers_pipeline = pipeline(
+            task=self.task,
+            model=self.model_pipeline,
+            tokenizer=self.tokenizer_pipeline,
+            device=(
+                self.device if isinstance(
+                    self.device, str) else (
+                    self.device.value)
+            ),
+            ** kwargs
+        )
+
+    def load_model(self) -> bool:
+        """
+        Load this model on the given device,
+
+        Returns:
+            bool: True if the model is successfully loaded.
+        """
+
+        if self.loaded:
+            return True
+        if self.device == Devices.RESET.value:
+            return False
+
+        # When the device is turn to meta, we must recreate them to load it
+        if (self.transformers_pipeline.device ==
+                torch.device(Devices.RESET.value)):
+            self.model_pipeline = self.model_class.from_pretrained(
+                self.model_path,
+                **self.model_pipeline_args
+            )
+
+        self.transformers_pipeline.model.to(device=(
+            self.device if isinstance(
+                self.device, str) else (
+                self.device.value)))
+        self.loaded = True
+        return True
+
+    def unload_model(self) -> bool:
+        """
+        Unloads the model
+        :return: True if the model is successfully unloaded
+        """
+        if not self.loaded:
+            return False
+
+        self.transformers_pipeline.model.to(Devices.RESET.value)
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+        self.loaded = False
+        return True
+
+    def generate_prompt(
+            self, prompt: Any,
+            **kwargs) -> Any:
+        """
+        Generates the prompt with the given option.
+
+        Args:
+            prompt (Any): The prompt.
+
+        Returns:
+            Any: Generated prompt.
+        """
+        return self.transformers_pipeline(prompt)
diff --git a/sdk/models/models_management.py b/sdk/models/models_management.py
index 956ea9e..0c51b30 100644
--- a/sdk/models/models_management.py
+++ b/sdk/models/models_management.py
@@ -1,6 +1,5 @@
 from typing import Optional, Dict
 from sdk.models import Model
-from sdk.options import Options
 
 
 class ModelsManagement:
@@ -16,13 +15,11 @@ def __init__(self):
         """
         self.loaded_model: Optional[Model] = None
         self.loaded_models_cache: Dict[str, Model] = {}
-        self.options_models: Dict[str, Options] = {}
 
-    def add_model(self, new_model: Model, model_options: Options) -> bool:
+    def add_model(self, new_model: Model) -> bool:
         """
         Adds a new model and his options to the management.
         :param new_model: The new model to add
-        :param model_options: The options of the new model to add
         :return: True if the model is successfully added
         """
         if new_model.model_name in self.loaded_models_cache:
@@ -30,7 +27,6 @@ def add_model(self, new_model: Model, model_options: Options) -> bool:
             return False
 
         self.loaded_models_cache[new_model.model_name] = new_model
-        self.options_models[new_model.model_name] = model_options
         return True
 
     def load_model(self, model_name: str) -> bool:
@@ -49,8 +45,7 @@ def load_model(self, model_name: str) -> bool:
             return False
 
         self.loaded_model = self.loaded_models_cache[model_name]
-        if not self.loaded_model.load_model(option=(
-                self.options_models[model_name])):
+        if not self.loaded_model.load_model():
             print("Something went wrong while unloading the model.")
             self.loaded_model = None
             return False
@@ -72,66 +67,29 @@ def unload_model(self) -> bool:
         self.loaded_model = None
         return True
 
-    def get_model_options(self, model_name: str) -> Options:
-        """
-        Gets the options of the model with the given name
-        :param model_name: The name of a model
-        :return: The object options of the model
-        """
-        return self.options_models[model_name]
-
-    def set_model_options(self, model_name: str, options: Options):
-        """
-        Set the options of the model with the given name
-        :param model_name: The name of a model
-        :param options: The object options of the model
-        """
-        self.options_models[model_name] = options
-
-    def generate_prompt(self, prompt: Optional[str] = None, **kwargs):
+    def generate_prompt(self, prompt: str,
+                        model_name: Optional[str] = None, **kwargs):
         """
         Generates the prompt for the loaded model with his stored options
-        :param prompt: The prompt to generate (if the prompt is empty, the
-            options.prompt will be used)
-        :param kwargs: more parameters to pass to the prompt generator
-        :return: The object of type link with the model category
-        """
-        if not self.loaded_model:
-            print("No model loaded. Load a model before generating prompts.")
-            return
-
-        return (
-            self.loaded_model.generate_prompt(
-                prompt,
-                self.options_models[
-                    self.loaded_model.model_name],
-                **kwargs
-            )
-        )
-
-    def generate_prompt_with_model_switch(self, prompt: Optional[str],
-                                          model_name: str, **kwargs):
-        """
-        Generates the prompt on the given model.
-        :param prompt: The prompt to generate (if the prompt is empty, the
-            options.prompt will be used)
-        :param model_name: The model to use for the prompt (if the model is not
-            loaded or different to the loaded model, the model will be loaded
-            before generating prompts
+        :param model_name: (Optional): the model name to load
+        :param prompt: The prompt to generate
         :param kwargs: more parameters to pass to the prompt generator
         :return: The object of type link with the model category
         """
-        if self.loaded_model.model_name != model_name:
-            self.unload_model()
+        if model_name:
+            if self.loaded_model.model_name != model_name:
+                self.unload_model()
 
         if not self.loaded_model:
+            if model_name:
+                print("No model loaded to generate.")
+                return
+
             self.load_model(model_name=model_name)
 
         return (
             self.loaded_model.generate_prompt(
                 prompt,
-                self.options_models[
-                    self.loaded_model.model_name],
                 **kwargs
             )
         )
diff --git a/sdk/options/__init__.py b/sdk/options/__init__.py
index 2fd5a97..7dbbcdb 100644
--- a/sdk/options/__init__.py
+++ b/sdk/options/__init__.py
@@ -1,4 +1,2 @@
 # flake8: noqa F401
-from sdk.options.options import Options, Devices
-from sdk.options.options_text_to_image import OptionsTextToImage
-from sdk.options.options_text_conversation import OptionsTextConversation
+from sdk.options.options import Devices
diff --git a/sdk/options/options.py b/sdk/options/options.py
index 8e29324..3c31617 100644
--- a/sdk/options/options.py
+++ b/sdk/options/options.py
@@ -1,6 +1,4 @@
-from abc import ABC
 from enum import Enum
-from typing import Union
 
 
 class Devices(Enum):
@@ -10,17 +8,3 @@ class Devices(Enum):
     GPU = "cuda"
     CPU = "cpu"
     RESET = "meta"
-
-
-class Options(ABC):
-    """
-    Abstract class defining the options used for models
-    """
-    device: Union[str, Devices]
-
-    def __init__(self, device: Devices):
-        """
-        Initializes the options class with the given device
-        :param device: The device to use generate prompt
-        """
-        self.device = device
diff --git a/sdk/options/options_text_conversation.py b/sdk/options/options_text_conversation.py
deleted file mode 100644
index c225942..0000000
--- a/sdk/options/options_text_conversation.py
+++ /dev/null
@@ -1,114 +0,0 @@
-import torch
-from sdk.options import Options, Devices
-from typing import Optional, Union, Dict, Any
-from transformers import ModelCard
-
-
-class OptionsTextConversation(Options):
-    """
-    Options for text-Generation models
-    """
-    """
-    Initializes the OptionsTextConversation.
-
-    Args:
-        prompt (str): The prompt or starting text for text generation.
-        model_name (str): The name of the model to use for text generation.
-        create_new_conv (bool): Whether to create a new conversation.
-            Default is False.
-        create_new_tokenizer (bool): Whether to create a new tokenizer.
-            Default is False.
-        delete_conv (bool): Whether to delete a conversation.
-            Default is False.
-        delete_tokenizer (bool): Whether to delete a tokenizer.
-            Default is False.
-        model_card (Optional[Union[str, ModelCard]]): The model card
-            providing details about the model.
-        framework (Optional[str]): The deep learning framework used for
-            the model, such as 'torch' or 'tf'.
-        task (str): The task for which the model is being used.
-            Default is an empty string.
-        num_workers (int): The number of worker processes for data loading.
-            Default is 8.
-        batch_size (int): The batch size for inference. Default is 1.
-        device (Union[str, Devices]): The device to use for inference,
-            specified from the `Devices` enum.
-        arg_parser (Optional[Dict[str, Any]]): Optional additional
-            arguments for the model.
-        torch_dtype (Optional[Union[str, torch.dtype]]): The data type for
-            PyTorch tensors, such as 'float32' or 'float64'.
-        binary_output (bool): Whether the output should be binary or text.
-            Default is False.
-        min_length_for_response (int): The minimum length of response
-            generated by the model. Default is 32.
-        minimum_tokens (int): The minimum number of tokens required for
-            a valid response. Default is 10.
-        attention_mask (bool): Whether to use attention mask.
-            Default is True.
-        pad_token_id (int): The ID of the padding token. Default is 50256.
-        eos_token_id (int): The ID of the end-of-sequence token.
-            Default is 50256.
-        tokenizer_id_to_use (int): The ID of the tokenizer to use.
-            Default is 0.
-        chat_id_to_use (int): The ID of the conversation to use.
-            Default is 0.
-    """
-
-    def __init__(self,
-                 prompt: str,
-                 model_name: str,
-                 create_new_conv: bool = False,
-                 delete_conv: bool = False,
-                 delete_tokenizer: bool = False,
-                 model_card: Optional[Union[str, ModelCard]] = None,
-                 framework: Optional[str] = None,
-                 task: str = "",
-                 num_workers: int = 8,
-                 batch_size: int = 1,
-                 device: Union[str, Devices] = -1,
-                 arg_parser: Optional[Dict[str, Any]] = None,
-                 torch_dtype: Optional[Union[str, torch.dtype]] = None,
-                 binary_output: bool = False,
-                 min_length_for_response: int = 32,
-                 minimum_tokens: int = 10,
-                 attention_mask: bool = True,
-                 pad_token_id: int = 50256,
-                 eos_token_id: int = 50256,
-                 tokenizer_id_to_use: int = 0,
-                 chat_id_to_use: int = 0,
-                 trust_remote_code: bool = False,
-                 max_length: Optional[int] = None,
-                 max_new_tokens: Optional[int] = None
-                 ):
-
-        super().__init__(device)
-        self.prompt = prompt
-        self.task = task
-        self.model_name = model_name
-        if model_card:
-            self.model_card = model_card
-        self.num_workers = num_workers
-        self.batch_size = batch_size
-        if framework:
-            self.framework = framework
-        if arg_parser:
-            self.arg_parser = arg_parser
-        if torch_dtype:
-            self.torch_dtype = torch_dtype
-        self.torch_dtype = torch_dtype
-        self.binary_output = binary_output
-        self.min_length_for_response = min_length_for_response
-        self.minimum_tokens = minimum_tokens
-        self.attention_mask = attention_mask
-        self.pad_token_id = pad_token_id
-        self.eos_token_id = eos_token_id
-        self.create_new_conv = create_new_conv
-        self.delete_tokenizer = delete_tokenizer
-        self.delete_conv = delete_conv
-        self.tokenizer_id_to_use = tokenizer_id_to_use
-        self.chat_id_to_use = chat_id_to_use
-        self.trust_remote_code = trust_remote_code
-        if max_length:
-            self.max_length = max_length
-        if max_new_tokens:
-            self.max_new_tokens = max_new_tokens
diff --git a/sdk/options/options_text_to_image.py b/sdk/options/options_text_to_image.py
deleted file mode 100644
index 3d77e9d..0000000
--- a/sdk/options/options_text_to_image.py
+++ /dev/null
@@ -1,264 +0,0 @@
-from typing import Optional, Union, List, Dict, Tuple, Any, Callable
-from diffusers.image_processor import PipelineImageInput
-from sdk.options import Options, Devices
-import torch
-
-
-class OptionsTextToImage(Options):
-    """
-    Options for text-to-image models
-    """
-    prompt: Union[str, List[str]]
-    prompt_2: Optional[Union[str, List[str]]] = None
-    image_width: Optional[int] = None
-    image_height: Optional[int] = None
-    num_inference_steps: int = 50
-    timesteps: List[int] = None
-    denoising_end: Optional[float] = None
-    guidance_scale: float = 5.0
-    negative_prompt: Optional[Union[str, List[str]]] = None
-    negative_prompt_2: Optional[Union[str, List[str]]] = None
-    num_images_per_prompt: Optional[int] = 1
-    eta: float = 0.0
-    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None
-    latents: Optional[torch.FloatTensor] = None
-    prompt_embeds: Optional[torch.FloatTensor] = None
-    negative_prompt_embeds: Optional[torch.FloatTensor] = None
-    pooled_prompt_embeds: Optional[torch.FloatTensor] = None
-    negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None
-    ip_adapter_image: Optional[PipelineImageInput] = None
-    output_type: Optional[str] = "pil"
-    return_dict: bool = True
-    cross_attention_kwargs: Optional[Dict[str, Any]] = None
-    guidance_rescale: float = 0.0
-    original_size: Optional[Tuple[int, int]] = None
-    crops_coords_top_left: Tuple[int, int] = (0, 0)
-    target_size: Optional[Tuple[int, int]] = None
-    negative_original_size: Optional[Tuple[int, int]] = None
-    negative_crops_coords_top_left: Tuple[int, int] = (0, 0)
-    negative_target_size: Optional[Tuple[int, int]] = None
-    clip_skip: Optional[int] = None
-    callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None
-    callback_on_step_end_tensor_inputs: List[str] = ["latents"]
-
-    def __init__(
-            self,
-            device: Union[str, Devices],
-            prompt: Union[str, List[str]],
-            prompt_2: Optional[Union[str, List[str]]] = None,
-            image_width: Optional[int] = None,
-            image_height: Optional[int] = None,
-            num_inference_steps: Optional[int] = None,
-            timesteps: List[int] = None,
-            denoising_end: Optional[float] = None,
-            guidance_scale: Optional[float] = None,
-            negative_prompt: Optional[Union[str, List[str]]] = None,
-            negative_prompt_2: Optional[Union[str, List[str]]] = None,
-            num_images_per_prompt: Optional[int] = None,
-            eta: Optional[float] = None,
-            generator: Optional[
-                Union[torch.Generator, List[torch.Generator]]] = None,
-            latents: Optional[torch.FloatTensor] = None,
-            prompt_embeds: Optional[torch.FloatTensor] = None,
-            negative_prompt_embeds: Optional[torch.FloatTensor] = None,
-            pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-            negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-            ip_adapter_image: Optional[PipelineImageInput] = None,
-            output_type: Optional[str] = None,
-            return_dict: Optional[bool] = None,
-            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
-            guidance_rescale: Optional[float] = None,
-            original_size: Optional[Tuple[int, int]] = None,
-            crops_coords_top_left: Tuple[int, int] = None,
-            target_size: Optional[Tuple[int, int]] = None,
-            negative_original_size: Optional[Tuple[int, int]] = None,
-            negative_crops_coords_top_left: Tuple[int, int] = None,
-            negative_target_size: Optional[Tuple[int, int]] = None,
-            clip_skip: Optional[int] = None,
-            callback_on_step_end: Optional[
-                Callable[[int, int, Dict], None]] = None,
-            callback_on_step_end_tensor_inputs: List[str] = None
-    ):
-        """
-        Initializes the OptionsTextToImage (Comment from stable diffusion xl)
-        :param device: The device to use generate prompt
-        :param prompt: The prompt to give to the model
-        :param image_width: The width of the resulting image (1024 by default)
-        :param image_height: The height of the resulting image
-            (1024 by default)
-        :param num_inference_steps (`int`, *optional*, defaults to 50):
-            The number of denoising steps. More denoising steps usually lead to
-             a higher quality image at the expense of slower inference.
-        :param timesteps (`List[int]`, *optional*):
-            Custom timesteps to use for the denoising process with schedulers
-            which support a `timesteps` argument in their `set_timesteps`
-            method. If not defined, the default behavior when
-            `num_inference_steps` is passed will be used.
-            Must be in descending order.
-        :param denoising_end (`float`, *optional*):
-            When specified, determines the fraction (between 0.0 and 1.0) of
-            the total denoising process to be completed before it is
-            intentionally prematurely terminated.
-            As a result, the returned sample will still retain a substantial
-            amount of noise as determined by the discrete timesteps selected by
-            the scheduler. The denoising_end parameter should ideally be
-            utilized when this pipeline forms a part of a
-            "Mixture of Denoisers" multi-pipeline setup.
-        :param guidance_scale (`float`, *optional*, defaults to 5.0):
-            Guidance scale as defined in Classifier-Free Diffusion Guidance.
-            `guidance_scale` is defined as `w` of equation 2. of the Imagen
-            Paper. Guidance scale is enabled by setting `guidance_scale > 1`.
-            Higher guidance scale encourages generating images closely linked
-            to the text `prompt`, usually at the expense of lower image quality
-        :param negative_prompt (`str` or `List[str]`, *optional*):
-            The prompt or prompts not to guide the image generation. If not
-            defined, one has to pass `negative_prompt_embeds` instead.
-            Ignored when not using guidance (i.e., ignored if `guidance_scale`
-            is less than `1`).
-        :param negative_prompt_2 (`str` or `List[str]`, *optional*):
-            The prompt or prompts not to guide the image generation to be sent
-            to `tokenizer_2` and `text_encoder_2`. If not defined,
-            `negative_prompt` is used in both text-encoders.
-        :param num_images_per_prompt (`int`, *optional*, defaults to 1):
-            The number of images to generate per prompt.
-        :param eta (`float`, *optional*, defaults to 0.0):
-            Corresponds to parameter eta (η) in the DDIM paper. Only applies to
-             [`schedulers.DDIMScheduler`], will be ignored for others.
-        :param generator (`torch.Generator` or `List[torch.Generator]`,
-            *optional*):
-            One or a list of torch generators to make generation deterministic.
-        :param latents (`torch.FloatTensor`, *optional*):
-            Pre-generated noisy latents, sampled from a Gaussian distribution,
-            to be used as inputs for image generation. Can be used to tweak the
-             same generation with different prompts. If not provided, a latents
-            tensor will be generated by sampling using the supplied random
-            generator.
-        :param prompt_embeds (`torch.FloatTensor`, *optional*):
-            Pre-generated text embeddings. Can be used to easily tweak text
-            inputs, *e.g.* prompt weighting. If not provided, text embeddings
-            will be generated from `prompt` input argument.
-        :param negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-            Pre-generated negative text embeddings. Can be used to easily tweak
-            text inputs, *e.g.* prompt weighting. If not provided,
-            negative_prompt_embeds will be generated from `negative_prompt`
-            input argument.
-        :param pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
-            Pre-generated pooled text embeddings. Can be used to easily tweak
-            text inputs, *e.g.* prompt weighting. If not provided, pooled text
-            embeddings will be generated from `prompt` input argument.
-        :param negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
-            Pre-generated negative pooled text embeddings.
-            Can be used to easily tweak text inputs, *e.g.* prompt weighting.
-            If not provided, pooled negative_prompt_embeds will be generated
-            from `negative_prompt` input argument.
-        :param ip_adapter_image: (`PipelineImageInput`, *optional*):
-            Optional image input to work with IP Adapters.
-        :param output_type (`str`, *optional*, defaults to `"pil"`):
-            The output format of the generated image. Choose between
-            PIL.Image.Image or np.array.
-        :param return_dict (`bool`, *optional*, defaults to `True`):
-            Whether or not to return a StableDiffusionXLPipelineOutput instead
-             of a plain tuple.
-        :param cross_attention_kwargs (`dict`, *optional*):
-            A kwargs dictionary that if specified is passed along to the
-            `AttentionProcessor` as defined under `self.processor` in
-            diffusers.models.attention_processor.
-        :param guidance_rescale (`float`, *optional*, defaults to 0.0):
-            Guidance rescale factor proposed by Common Diffusion Noise
-            Schedules and Sample Steps are Flawed. `guidance_scale` is defined
-            as `φ` in equation 16. of Common Diffusion Noise Schedules and
-            Sample Steps are Flawed. Guidance rescale factor should fix
-            overexposure when using zero terminal SNR.
-        :param original_size
-            (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
-            If `original_size` is not the same as `target_size` the image will
-            appear to be down- or upsampled. `original_size` defaults to
-            `(height, width)` if not specified. Part of SDXL's
-            micro-conditioning as explained in section 2.2 of the documentation
-        :param crops_coords_top_left
-            (`Tuple[int]`, *optional*, defaults to (0, 0)):
-            `crops_coords_top_left` can be used to generate an image that
-            appears to be "cropped" from the position `crops_coords_top_left`
-            downwards. Favorable, well-centered images are usually achieved by
-            setting `crops_coords_top_left` to (0, 0). Part of SDXL's
-            micro-conditioning as explained in section 2.2 of the documentation
-        :param target_size
-            (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
-            For most cases, `target_size` should be set to the desired height
-            and width of the generated image. If not specified it will default
-            to `(height, width)`. Part of SDXL's micro-conditioning as
-            explained in section 2.2 of the documentation.
-        :param negative_original_size
-            (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
-            To negatively condition the generation process based on a specific
-            image resolution. Part of SDXL's micro-conditioning as explained in
-            section 2.2 of the documentation.
-        :param negative_crops_coords_top_left
-            (`Tuple[int]`, *optional*, defaults to (0, 0)):
-            To negatively condition the generation process based on specific
-            crop coordinates. Part of SDXL's micro-conditioning as explained in
-             section 2.2 of the documentation.
-        :param negative_target_size
-            (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
-            To negatively condition the generation process based on a target
-            image resolution. It should be as same as the `target_size` for
-            most cases. Part of SDXL's micro-conditioning as explained in
-            section 2.2 of the documentation.
-        :param callback_on_step_end (`Callable`, *optional*):
-            A function that calls at the end of each denoising step during the
-            inference. The function is called with the following arguments:
-            `callback_on_step_end(self: DiffusionPipeline, step: int, timestep:
-            int, callback_kwargs: Dict)`. `callback_kwargs` will include a list
-            of all tensors as specified by `callback_on_step_end_tensor_inputs`
-        :param callback_on_step_end_tensor_inputs (`List`, *optional*):
-            The list of tensor inputs for the `callback_on_step_end` function.
-            The tensors specified in the list will be passed as
-            `callback_kwargs` argument. You will only be able to include
-            variables listed in the `._callback_tensor_inputs` attribute of
-            your pipeline class.
-        """
-        super().__init__(device)
-        self.prompt = prompt
-        self.prompt_2 = prompt_2
-        self.image_width = image_width
-        self.image_height = image_height
-        if num_inference_steps:
-            self.num_inference_steps = num_inference_steps
-        self.timesteps = timesteps
-        self.denoising_end = denoising_end
-        if guidance_scale:
-            self.guidance_scale = guidance_scale
-        self.negative_prompt = negative_prompt
-        self.negative_prompt_2 = negative_prompt_2
-        if num_images_per_prompt:
-            self.num_images_per_prompt = num_images_per_prompt
-        if eta:
-            self.eta = eta
-        self.generator = generator
-        self.latents = latents
-        self.prompt_embeds = prompt_embeds
-        self.negative_prompt_embeds = negative_prompt_embeds
-        self.pooled_prompt_embeds = pooled_prompt_embeds
-        self.negative_pooled_prompt_embeds = negative_pooled_prompt_embeds
-        self.ip_adapter_image = ip_adapter_image
-        if output_type:
-            self.output_type = output_type
-        if return_dict:
-            self.return_dict = return_dict
-        self.cross_attention_kwargs = cross_attention_kwargs
-        if guidance_rescale:
-            self.guidance_rescale = guidance_rescale
-        self.original_size = original_size
-        if crops_coords_top_left:
-            self.crops_coords_top_left = crops_coords_top_left
-        self.target_size = target_size
-        self.negative_original_size = negative_original_size
-        if negative_crops_coords_top_left:
-            self.negative_crops_coords_top_left = (
-                                                negative_crops_coords_top_left)
-        self.negative_target_size = negative_target_size
-        self.clip_skip = clip_skip
-        self.callback_on_step_end = callback_on_step_end
-        if callback_on_step_end_tensor_inputs:
-            self.callback_on_step_end_tensor_inputs = (
-                                            callback_on_step_end_tensor_inputs)
diff --git a/sdk/options/options_tokenizer.py b/sdk/options/options_tokenizer.py
deleted file mode 100644
index 78260a5..0000000
--- a/sdk/options/options_tokenizer.py
+++ /dev/null
@@ -1,33 +0,0 @@
-from typing import Union, Optional
-from sdk import Devices
-
-
-class OptionsTokenizer:
-    """
-    Abstract class defining the options used for models
-    """
-    device: Union[str, Devices]
-    padding_side: str = "left"
-    return_tensors: Optional[str] = None
-
-    def __init__(self, device: Union[str, Devices],
-                 padding_side: Optional[str],
-                 return_tensors: Optional[str]):
-        """
-        Initializes the options class with the given device and other options.
-
-        Args:
-            device (Union[str, Devices]): The device used for generation.
-            padding_side (str): The padding side,
-                either 'left' or 'right'. Defaults to 'left'.
-            return_tensors (str): The return tensor format. Defaults to 'pt'.
-        """
-        self.device = device
-        if padding_side:
-            if padding_side not in ['left', 'right']:
-                raise ValueError(
-                    "padding_side must be either 'left' or 'right'")
-        else:
-            self.padding_side = padding_side
-        if return_tensors:
-            self.return_tensors = return_tensors
diff --git a/sdk/tokenizers/tokenizer.py b/sdk/tokenizers/tokenizer.py
index 3b8b690..bae91c2 100644
--- a/sdk/tokenizers/tokenizer.py
+++ b/sdk/tokenizers/tokenizer.py
@@ -1,79 +1,57 @@
-import torch
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, PreTrainedTokenizer
+from typing import Union
 
-from sdk.options.options_tokenizer import OptionsTokenizer
+from sdk.options.options import Devices
 
 
 class Tokenizer:
     """
     Abstract base class for all tokenizers
     """
-    pipeline: AutoTokenizer
-    options: OptionsTokenizer
-    model_name: str
+    pipeline: PreTrainedTokenizer
+    device: Union[str, Devices]
     tokenizer_name: str
     tokenizer_path: str
-    user_input: str
-    chatbot_output: str
-    conversation_active: bool = False
-    return_tensors = 'pt'
 
-    def __init__(self, model_name: str,
-                 tokenizer_path: str,
+    def __init__(self,
                  tokenizer_name: str,
-                 options: OptionsTokenizer):
+                 tokenizer_path: str,
+                 device: Union[str, Devices],
+                 **kwargs):
         """
         Initializes the TokenizerObject class with the given parameters.
 
         Args:
-            model_name (str): The name of
-                the model associated with the tokenizer .
             tokenizer_name (str): The name of the tokenizer
             tokenizer_path (str): The path of the tokenizer
-            options (OptionsTokenizer): The options for the tokenizer.
         """
-        self.model_name = model_name
         self.tokenizer_name = tokenizer_name
         self.tokenizer_path = tokenizer_path
-        self.options = options
-        self.create_tokenizer()
+        self.device = device
+        self.create_tokenizer(**kwargs)
 
-    def create_tokenizer(self):
+    def create_tokenizer(self, **kwargs):
         """
         Creates the tokenizer to use
         """
         self.pipeline = AutoTokenizer.from_pretrained(
             self.tokenizer_path,
-            trust_remote_code=True,
-            padding_side=self.options.padding_side
+            **kwargs
         )
 
-    def prepare_input(self, prompt: str, history: str):
+    def encode(self, prompt: str, **kwargs):
         """
         Tokenizes the given prompt and prepares it for model input.
 
-        Args:
-            history: chat history
-            prompt (str): The input prompt.
+        :param prompt: (str): The prompt to encode
+        :param kwargs: (dict): Additional options to pass to the encode method
 
         Returns:
             torch.Tensor: The tokenized and formatted input tensor.
         """
         return self.pipeline.encode(
-            history,
             prompt,
-            return_tensors=self.return_tensors,
-            truncation=True
-        ).to(self.options.device)
-
-    def decode_model_output(self, outputs: torch.Tensor):
-        """
-        Decodes the model output tensor to obtain the generated text.
-
-        Args:
-            outputs (torch.Tensor): The model output tensor.
-
-        Returns:
-            str: The decoded generated text.
-        """
-        return self.pipeline.decode(outputs[0]).strip()
+            **kwargs,
+        ).to(device=(self.device if isinstance(
+            self.device, str) else (
+                self.device.value)))

From 69eb832fad761a4a1162d8b1c272ba35721ed975 Mon Sep 17 00:00:00 2001
From: Quentin <63395737+SkelNeXus@users.noreply.github.com>
Date: Wed, 20 Mar 2024 16:42:29 +0100
Subject: [PATCH 3/3] 83 sdk add demo with text generation (#84)

* demo for text generation

* Fix on model transformers and demo text gen
---
 main.py                          | 10 ++++++---
 sdk/demo/__init__.py             |  1 +
 sdk/demo/demo_text_generation.py | 37 ++++++++++++++++++++++++++++++++
 sdk/models/model_transformers.py |  2 +-
 4 files changed, 46 insertions(+), 4 deletions(-)
 create mode 100644 sdk/demo/demo_text_generation.py

diff --git a/main.py b/main.py
index 7f8850a..a2e4712 100644
--- a/main.py
+++ b/main.py
@@ -1,21 +1,25 @@
 import argparse
-from sdk.demo import DemoTextConv, DemoTextToImg
+from sdk.demo import DemoTextConv, DemoTextToImg, DemoTextGen
 
 
 if __name__ == '__main__':
 
     parser = argparse.ArgumentParser(description='Choose a Model type :'
-                                                 ' TextConv or TextToImg')
+                                                 ' TextConv, TextToImg '
+                                                 'or TextGen')
     subparser = parser.add_subparsers(dest='option')
 
     conv = subparser.add_parser('TextConv')
     img = subparser.add_parser('TextToImg')
+    gen = subparser.add_parser('TextGen')
     args = parser.parse_args()
 
     match args.option:
         case 'TextConv':
             DemoTextConv()
-        case ('TextToImg'):
+        case 'TextToImg':
             DemoTextToImg()
+        case 'TextGen':
+            DemoTextGen()
         case _:
             DemoTextToImg()
diff --git a/sdk/demo/__init__.py b/sdk/demo/__init__.py
index 3b91259..c5f7a15 100644
--- a/sdk/demo/__init__.py
+++ b/sdk/demo/__init__.py
@@ -1,3 +1,4 @@
 # flake8: noqa F401
 from .demo_text_conv import DemoTextConv
 from .demo_text_to_img import DemoTextToImg
+from .demo_text_generation import DemoTextGen
diff --git a/sdk/demo/demo_text_generation.py b/sdk/demo/demo_text_generation.py
new file mode 100644
index 0000000..1e7d30b
--- /dev/null
+++ b/sdk/demo/demo_text_generation.py
@@ -0,0 +1,37 @@
+from sdk.options import Devices
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+from sdk.models import ModelTransformers
+
+
+class DemoTextGen:
+    """
+    This class demonstrates a text conversation using a chatbot model.
+    """
+
+    def __init__(self):
+        """
+        Initializes the DemoTextConv class with predefined options and models.
+        """
+        model_path = "microsoft/phi-2"
+        tokenizer_path = "microsoft/phi-2"
+
+        model_transformers = ModelTransformers(
+            model_name="model",
+            model_path=model_path,
+            tokenizer_path=tokenizer_path,
+            task="text-generation",
+            model_class=AutoModelForCausalLM,
+            tokenizer_class=AutoTokenizer,
+            device=Devices.GPU
+        )
+
+        model_transformers.load_model()
+
+        result = model_transformers.generate_prompt(
+            prompt="I'm looking for a movie - what's your favourite one?",
+            max_length=300,
+            truncation=True
+        )
+
+        print(result)
diff --git a/sdk/models/model_transformers.py b/sdk/models/model_transformers.py
index 62f7d01..6ff846a 100644
--- a/sdk/models/model_transformers.py
+++ b/sdk/models/model_transformers.py
@@ -143,4 +143,4 @@ def generate_prompt(
         Returns:
             Any: Generated prompt.
         """
-        return self.transformers_pipeline(prompt)
+        return self.transformers_pipeline(prompt, **kwargs)