Repository: hand-e-fr/OpenHosta (language: python)

Issue:
Classification gestion for encoder in predict 

Feature/normalization



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 3c193e2..abd4310 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,5 @@
 # OpenHosta 
-v2.0.2 - Opensource Project
+v2.1.0rc-1 - Opensource Project
 
 **- The future of development is human -**
 
@@ -79,7 +79,7 @@ For further questions or assistance, please refer to partner hand-e or contact u
 **Authors**:
    - Emmanuel Batt: Manager and Coordinator, Founder of Hand-e
    - William Jolivet: DevOps, SysAdmin
-   - Léandre Ramos: MLOps, IA developer
+   - Léandre Ramos: IA developer
    - Merlin Devillard: UX designer, Product Owner
 
 GitHub: https://github.com/hand-e-fr/OpenHosta
diff --git a/docs/doc.md b/docs/doc.md
index e698170..c8efa53 100644
--- a/docs/doc.md
+++ b/docs/doc.md
@@ -79,8 +79,16 @@ Let's **get started**! First here's the **table of contents** to help you naviga
     - [Body Functions](#body-functions)
       - [`Example`](#example)
       - [`Thought`](#thought)
+  - [`predict` Function](#predict-function)
+    - [PredictConfig](#predictconfig-class)
   - [`thinkof` Function](#thinkof-function)
   - [`ask` function](#ask-function)
+  - [`generate_data` func](#generate_data-func)
+    - [Parameters](#parameters)
+    - [Returns](#returns)
+    - [Raises](#raises)
+    - [Example](#example)
+    - [How It Works](#how-it-works)
   - [Advanced configuration](#advanced-configuration)
     - [Models](#models)
       - [Inheriting from the Model Class](#inheriting-from-the-model-class)
@@ -182,7 +190,7 @@ from typing import List
 
 def example_function(a:int, b:str)->List[str]:
 ```
-  - **The doctring** is the other key element. This is where you describe the behavior of the function. Be precise and concise. Describe the input parameters and the nature of the output, as in a docstring. Feel free to try out lots of things, prompt engineering is not a closed science. :)
+  - **The docstring** is the other key element. This is where you describe the behavior of the function. Be precise and concise. Describe the input parameters and the nature of the output, as in a docstring. Feel free to try out lots of things, prompt engineering is not a closed science. :)
 
 ```python
 my_model = config.Model(
@@ -205,7 +213,7 @@ def find_name_age(sentence:str, id:dict)->dict:
     """
     return emulate(model=my_model)
 
-ret = find_name_age("l'âge du capitaine est d'un an de plus que celui du mousse qui a lui 22 ans", {"capitaine": 0, "mousse": 0})
+ret = find_name_age("the captain's age is one year more than the cabin boy's, who is 22 years old", {"captain": 0, "cabin boy": 0})
 print(ret)
 ```
 
@@ -353,6 +361,141 @@ print(car_advice("I would two buy a new car, but I would like an electric becaus
 # [{'task': 'identify the context and the need of the user'}, {'task': 'Look at the car available to find a car matching his needs'}, {'task': 'Return the most relevant car, if no car is matching return None'}, {'task': 'identify the context and the need of the user'}, {'task': 'Look at the car available to find a car matching his needs'}, {'task': 'Return the most relevant car, if no car is matching return None'}]
 ```
 
+## `predict` Function
+
+The `predict` function is the second main feature of OpenHosta ! This function allows you to create **specific neural networks** based on the specifications you provide. Here's a breakdown to help you understand it:
+
+
+
+The `predict` function can be used in function or class method by simply returns it. Its primary goal is to create a model tailored to the function it is called in. Currently, it supports two model types:
+
+- **Linear Regression**: For prediction tasks by simply returns an `int`or a `float` :
+  ```python
+  from Openhosta import predict, config
+
+  config.set_default_apiKey("put-your-api-key-here")
+  
+  def example_linear_regression(years : int, profession : str) -> float:
+      """
+      this function predict the salary based on the profession years.
+      """
+      return predict(verbose=2)
+  
+  print(example_linear_regression(1, "engineer"))
+  ```
+- **Classification**: For classifying multiple values among predefined categories in a `Literal` from the typing module :
+  ```python
+  from typing import Literal
+  from Openhosta import predict, config
+
+  config.set_default_apiKey("put-your-api-key-here")
+
+  output = Literal["Good", "Bad"]
+  
+  def example_classification(word: str) -> output:
+      """
+      this function detects if a words is good or bad
+      """
+      return predict(verbose=2)
+  
+  print(example_classification("Bad"))
+  ```
+
+
+Additionally like you can see, `predict` can generate a dataset if none is provided in the [PredictConfig](#predictconfig-class), allowing users to see how a large language model (LLM) understands the problem and generates relevant data. By default, the data generation uses GPT-4o by OpenAI, the same oracle used in the [emulate](#emulate-function) function .
+
+
+
+### Parameters
+The `predict` function supports the following parameters:
+
+- `verbose`: Controls the level of output information:
+  - `0`: No output.
+  - `1`: Basic output (default).
+  - `2`: Detailed output.
+
+
+
+- `oracle`: Specifies the model used for data generation. If set to `None`, no model will be used to handle missing data generation.
+
+- `config`: Accepts a `Predictconfig` object for advanced configuration of model creation.
+
+## `PredictConfig` class
+
+The `PredictConfig` class provides advanced options for configuring the creation and management of *predict* models. Here’s a detailed breakdown:
+
+```python
+from Openhosta import PredictConfig
+
+model_config = PredictConfig(
+    name="model_test",
+    path="./__hostacache__",
+    complexity=5,
+    growth_rate=1.5,
+    coef_layers=100,
+    epochs=100,
+    batch_size=32,
+    max_tokens=1,
+    dataset_path="./path_to_dataset.csv",
+    generated_data=100,
+    normalize=False
+)
+```
+
+### Features
+
+#### **Path Management**
+- `path` (`str`): Specifies where data will be stored. Default: `./__hostacache__/`.
+- `name` (`str`): Sets the directory name for storing model-related information. Default: the name of the Hosta-injected function.
+
+#### **Architecture Configuration**
+- `complexity` (`int`): Adjusts the model's complexity by adding or removing layers. Default: `5`.
+- `growth_rate` (`float`): Determines the rate of increase in layer size. Default: `1.5`.
+- `coef_layers` (`int`): Defines the maximum possible layer size based on the highest value between inputs and outputs. Default: `100`.
+
+#### **Training Configuration**
+- `normalize` (`bool`) : Specify if the data for training need to be normalize between -1 and 1. Default False, only possible with **Linear Regression** models
+- `epochs` (`int`): Sets the number of training iterations. Default: calculated based on dataset size and batch size.
+- `batch_size` (`int`): Specifies the number of examples processed before weight updates. Default: 5% of the dataset size if possible.
+
+#### **Dataset Management**
+- `max_tokens` (`int`): Limits the number of words a `str` input can contain, as the model adapts neuron sizes accordingly. Default: `1`.
+  - **Warning**: Current model architectures do not perform well with natural language processing tasks. For such cases, use the *emulate* feature instead. NLP architecture support is coming soon.
+- `dataset_path` (`str`): Provides a custom dataset path. Default: `None`, 
+  - **Warning**: Only `csv` and `jsonl` files are supported for now. For `csv`, please set the prediction columns to `_outputs`. and for `jsonl` please set the last element
+- `generated_data` (`int`): Specifies the target number of data points for LLM generation (approximate). Default: `100`.
+
+---
+
+### Example Usage
+
+```python
+from OpenHosta import predict, PredictConfig
+
+# Configuring the model
+config_predict = PredictConfig(
+    path="./__hostacache__",
+    name="test_openhosta",
+    complexity=5,
+    growth_rate=1.5,
+    layers_coefficien=100,
+    epochs=45,
+    batch_size=64,
+    max_tokens=1,
+    dataset_path="./dataset.csv",
+    normalize=False
+)
+
+# Using the predict function with the configuration
+def demo_open_hosta(number: int, message: str) -> int:
+    """
+    this function is just here for an example :)
+    """
+    return predict(config=config_predict, oracle=None, verbose=2)
+
+print(demo_open_hosta(42, "Hello World!"))
+```
+
 ## `thinkof` Function
 
 **Lambda** functions in Python provide a way to create small, anonymous functions. These are defined using the lambda keyword and can have any number of input parameters but only a single expression.
@@ -433,6 +576,83 @@ As seen above takes 2 or more argument. The two first arguments are mandatory. `
 
 **Note** : ***this feature uses the default model.***
 
+## `generate_data` func
+
+Generate a dataset based on a given function and the number of samples. This function uses a synthetic data generator to create realistic input-output pairs for a given callable Python function based on its defined parameters, examples, and return type.
+
+### Parameters
+
+- **`func`** (`Callable`):  
+  The target function used to generate the dataset. This function must take specific inputs and return outputs to be used for creating the dataset.  
+  Proper type annotations and a clear docstring for the `func` are recommended to enhance the quality of generated data.
+
+- **`num_samples`** (`int`):  
+  The number of samples to generate. If the number exceeds 100, the function intelligently splits the data requests into manageable chunks.
+
+- **`oracle`** (`Optional[Model]`, Optional):  
+  The model or "oracle" used to assist with generating synthetic data.  
+  By default, the function uses the system's predefined default model.
+
+- **`verbose`** (`Union[Literal[0, 1, 2], bool]`, default=`2`):  
+  Defines the verbosity level for logging the data generation process:  
+  - `0` or `False`: No logging.
+  - `1`: Minimal logging.
+  - `2` or `True`: Detailed logging, providing insights during data generation.
+
+### Returns
+
+- **`HostaDataset`**:  
+  An instance of `HostaDataset`, representing the generated dataset. This dataset can be saved to disk (CSV, JSON, JSONL) or iterated over for input-output pairs.
+
+### Raises
+
+- **`TypeError`**:  
+  Raised if the provided `func` is not callable or lacks sufficient information to generate data (such as missing type annotations).
+
+### Example
+
+The following example demonstrates how to define a function, generate synthetic data using `generate_data`, and save the resulting dataset.
+
+```python
+from typing import Literal
+from OpenHosta import generate_data, HostaDataset, example, emulate, SourceType
+
+def detect_mood(message: str) -> Literal["positive", "negative", "neutral"]:
+    """
+    Analyze the mood conveyed in a text message.
+    """
+    # Provide pre-defined examples to guide synthetic data generation
+    example(message="I feel lonely...", hosta_out="negative")
+    example(message="I am happy!", hosta_out="positive")
+    example(message="I have a cat", hosta_out="neutral")
+    return emulate()
+
+# Generate a dataset with 50 examples
+dataset: HostaDataset = generate_data(detect_mood, 50)
+
+# Save the dataset as a CSV file
+dataset.save_data("detect_mood.csv", SourceType.CSV)
+
+# Print each input-output pair in the dataset
+correct = 0
+
+for data in dataset.data:
+    print(f"{data.input[0].strip('"')} expected {data.output} got {detect_mood(data.input[0].strip('"'))}")
+    if data.output == detect_mood(data.input[0].strip('"')):
+        correct += 1
+
+print(f"Accuracy: {correct}/{len(dataset.data)}, {correct/len(dataset.data)*100}%")
+```
+
+### How It Works
+
+- **Define the Function**:  
+  The target function (`detect_mood` in the example) must be well-defined, preferably with type annotations and examples to guide the data generation process.
+- **Generate Synthetic Data**:
+  Use `generate_data` to produce a dataset by specifying the number of samples and optionally overriding the default model with a custom `oracle`.
+- **Save or Process Dataset**:  
+  The returned dataset (`HostaDataset` instance) provides methods to save it in various formats (CSV, JSON, JSONL) or iterate over its contents for further analysis.
+
 ## Advanced configuration
 
 ### Models
@@ -522,8 +742,8 @@ class MyModel(Model):
         return l_ret
 
 new_model = MyModel(
-    model="model-name"
-    base_url="base-url"
+    model="model-name",
+    base_url="base-url",
     api_key="put-your-api-key-here"
 )
 
diff --git a/pyproject.toml b/pyproject.toml
index e75f9c8..a17d697 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,12 +18,27 @@ license = {file = "LICENSE"}
 requires-python = ">=3.8,<=3.13"
 classifiers = [
     "Development Status :: 5 - Production/Stable",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: Implementation :: CPython",
     "License :: OSI Approved :: MIT License",
     "Operating System :: OS Independent",
     "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Intended Audience :: Education",
     "Natural Language :: French",
-    "Topic :: Software Development :: Code Generators"
+    "Natural Language :: English",
+    "Topic :: Software Development :: Code Generators",
+    "Topic :: Scientific/Engineering",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries",
+    "Typing :: Typed"
 ]
 dependencies = [
     "requests>=2.32.3",
diff --git a/qodana.yaml b/qodana.yaml
new file mode 100644
index 0000000..84e3e49
--- /dev/null
+++ b/qodana.yaml
@@ -0,0 +1,29 @@
+#-------------------------------------------------------------------------------#
+#               Qodana analysis is configured by qodana.yaml file               #
+#             https://www.jetbrains.com/help/qodana/qodana-yaml.html            #
+#-------------------------------------------------------------------------------#
+version: "1.0"
+
+#Specify inspection profile for code analysis
+profile:
+  name: qodana.starter
+
+#Enable inspections
+#include:
+#  - name: <SomeEnabledInspectionId>
+
+#Disable inspections
+#exclude:
+#  - name: <SomeDisabledInspectionId>
+#    paths:
+#      - <path/where/not/run/inspection>
+
+#Execute shell command before Qodana execution (Applied in CI/CD pipeline)
+#bootstrap: sh ./prepare-qodana.sh
+
+#Install IDE plugins before Qodana execution (Applied in CI/CD pipeline)
+#plugins:
+#  - id: <plugin.id> #(plugin id can be found at https://plugins.jetbrains.com)
+
+#Specify Qodana linter for analysis (Applied in CI/CD pipeline)
+linter: jetbrains/qodana-python:latest
diff --git a/src/OpenHosta/OpenHosta.py b/src/OpenHosta/OpenHosta.py
index 81e861f..fbc6850 100644
--- a/src/OpenHosta/OpenHosta.py
+++ b/src/OpenHosta/OpenHosta.py
@@ -6,10 +6,15 @@
 from .exec.thought import thought
 from .exec.example import example
 from .exec.emulate import emulate
+from .exec.generate_data import generate_data
+from .exec.predict.dataset.dataset import HostaDataset, SourceType
 from .core import config
 from .core.config import Model, DefaultManager
 from .utils.meta_prompt import print_last_prompt
 
+from .exec.predict.predict import predict
+from .exec.predict.predict_config import PredictConfig
+
 import os
 
 DefaultManager.set_default_model(
@@ -22,8 +27,13 @@
     "emulate",
     "thought",
     "example",
+    "predict",
+    "PredictConfig",
     "thinkof",
     "ask",
     "EMULATE_PROMPT",
-    "print_last_prompt"
+    "print_last_prompt",
+    "generate_data",
+    "HostaDataset",
+    "SourceType"
 )
diff --git a/src/OpenHosta/core/analizer.py b/src/OpenHosta/core/analizer.py
index ca7b346..537d8a8 100644
--- a/src/OpenHosta/core/analizer.py
+++ b/src/OpenHosta/core/analizer.py
@@ -145,10 +145,10 @@ def func_call(self) -> Tuple[str, OrderedDict[str, Any]]:
     @property
     def func_type(self) -> Tuple[List[Any], Any]:
         """
-        Get the _inputs and _outputs types of the function.
+        Get the inputs and outputs types of the function.
 
         Returns:
-            A tuple containing the _inputs types and _outputs type.
+            A tuple containing the inputs types and outputs type.
         """
         input_types = [
             param.annotation for param in self.sig.parameters.values()
diff --git a/src/OpenHosta/core/checker.py b/src/OpenHosta/core/checker.py
index 06a11db..7778ea3 100644
--- a/src/OpenHosta/core/checker.py
+++ b/src/OpenHosta/core/checker.py
@@ -15,15 +15,15 @@
 
 class HostaChecker:
     """
-    A class used to check and convert the _outputs of a Language Model (LLM) to the type specified in a function's annotation.
+    A class used to check and convert the outputs of a Language Model (LLM) to the type specified in a function's annotation.
 
     Args:
-        func (Func): A function object that contains the type annotations for the LLM _outputs.
-        data (dict): A dictionary containing the LLM _outputs data to be checked and converted.
+        func (Func): A function object that contains the type annotations for the LLM outputs.
+        data (dict): A dictionary containing the LLM outputs data to be checked and converted.
 
     Attributes:
-        func (Func): The function object containing the type annotations for the LLM _outputs.
-        data (dict): The LLM _outputs data to be checked and converted.
+        func (Func): The function object containing the type annotations for the LLM outputs.
+        data (dict): The LLM outputs data to be checked and converted.
         checked (Any): The checked and converted data. If `data` contains a "return" key, its value is used as the checked data. Otherwise, `data` is used as the checked data.
         is_passed (bool): A flag indicating whether the checked data should be converted or not. It is set to True if `data` contains a "return" key.
     """
@@ -50,7 +50,7 @@ def _default(x: Any) -> Any:
         """
         return x
 
-    def convert(self, typ: Type[T]) -> Dict[Type[T], Optional[Callable[[Any], T]]]:
+    def convert(self, typ: Type[T]) -> Callable[[Any], T]:
         """
         A method to create a conversion function for a given type.
 
@@ -58,25 +58,23 @@ def convert(self, typ: Type[T]) -> Dict[Type[T], Optional[Callable[[Any], T]]]:
             typ (Type[T]): The type for which a conversion function needs to be created.
 
         Returns:
-            Dict[Type[T], Optional[Callable[[Any], T]]]: A dictionary mapping types to their corresponding conversion functions.
+            Callable[[Any], T]: A conversion function for the given type.
         """
-        convertMap = {
+        convert_map = {
             NoneType: lambda x: None,
-            str: lambda x: str(x),
-            int: lambda x: int(x),
-            float: lambda x: float(x),
-            list: lambda x: list(x),
-            set: lambda x: set(x),
-            frozenset: lambda x: frozenset(x),
-            tuple: lambda x: tuple(x),
-            bool: lambda x: bool(x),
-            dict: lambda x: dict(x),
-            complex: lambda x: complex(x),
+            str: str,
+            int: int,
+            float: float,
+            list: list,
+            set: set,
+            frozenset: frozenset,
+            tuple: tuple,
+            bool: bool,
+            dict: dict,
+            complex: complex,
             bytes: lambda x: bytes(x, encoding='utf-8') if isinstance(x, str) else bytes(x),
         }
-        if typ not in convertMap.keys():
-            return self._default.__func__
-        return convertMap[typ]
+        return convert_map.get(typ, self._default.__func__)
 
     def convert_annotated(self) -> Any:
         """
diff --git a/src/OpenHosta/core/config.py b/src/OpenHosta/core/config.py
index df47b6f..fa1563a 100644
--- a/src/OpenHosta/core/config.py
+++ b/src/OpenHosta/core/config.py
@@ -30,10 +30,11 @@ class Model:
 
     _SYS_PROMPT = ""
 
-    def __init__(self, model: str = None, base_url: str = None, api_key: str = None):
+    def __init__(self, model: str = None, base_url: str = None, api_key: str = None, timeout: int = 30):
         self.model = model
         self.base_url = base_url
         self.api_key = api_key
+        self.timeout = timeout
         self._last_request = None
         self._used_tokens = 0
         self._nb_requests = 0
@@ -86,7 +87,7 @@ def api_call(
         for key, value in llm_args.items():
             l_body[key] = value
         try:
-            response = requests.post(self.base_url, headers=headers, json=l_body, timeout=30)
+            response = requests.post(self.base_url, headers=headers, json=l_body, timeout=self.timeout)
 
             if response.status_code  != 200:
                 response_text = response.text
diff --git a/src/OpenHosta/core/hosta.py b/src/OpenHosta/core/hosta.py
index 37933b7..8450634 100644
--- a/src/OpenHosta/core/hosta.py
+++ b/src/OpenHosta/core/hosta.py
@@ -59,10 +59,10 @@ def __new__(cls, *args, **kwargs) -> 'Hosta':
                 "[Hosta.__new__] The function {} must be called in a function/method."
                 .format(cls._extend(back_level=2)[0].__name__)
             )
-        if (hasattr(cls._obj[0], "Hosta")):
+        if hasattr(cls._obj[0], "Hosta"):
             return cls._obj[0].Hosta
         instance = super().__new__(cls)
-        cls._attach(cls._obj[0], {"Hosta": instance})
+        cls.attach(cls._obj[0], {"Hosta": instance})
         return instance
 
     def __init__(self, *, caller_analysis: bool = True):
@@ -122,17 +122,16 @@ def _bdy_add(self, key: MemKey, value: MemValue) -> None:
             value (MemValue): The value to be stored in the memory node.
         """
         seen: List[MemKey] = []
-        previous: MemKey = None
 
         if self._infos.f_mem is None:
             self._infos.f_mem = []
-            id = 0
+            mem_id = 0
         else:
-            id = 0
+            mem_id = 0
             for node in self._infos.f_mem:
                 if node.key == key:
-                    id += 1
-        new = MemoryNode(key=key, id=id, value=value)
+                    mem_id += 1
+        new = MemoryNode(key=key, id=mem_id, value=value)
         self._infos.f_mem.append(new)
         previous = new
         for node in self._infos.f_mem:
diff --git a/src/OpenHosta/core/inspector.py b/src/OpenHosta/core/inspector.py
index ffc6819..bf84883 100644
--- a/src/OpenHosta/core/inspector.py
+++ b/src/OpenHosta/core/inspector.py
@@ -128,10 +128,10 @@ def _get_obj_from_func(
             raise FrameError(
                 "[HostaInspector._extend] The foud object isn't a callable.")
 
-        return (func, caller)
+        return func, caller
 
     @staticmethod
-    def _attach(obj: Callable, attr: Dict[str, Any]) -> Optional[bool]:
+    def attach(obj: Callable, attr: Dict[str, Any]) -> Optional[bool]:
         """
         Attaches attributes to a function or method.
 
diff --git a/src/OpenHosta/core/logger.py b/src/OpenHosta/core/logger.py
new file mode 100644
index 0000000..c79b6c6
--- /dev/null
+++ b/src/OpenHosta/core/logger.py
@@ -0,0 +1,108 @@
+from typing import Union, Literal, Optional
+from enum import Enum
+import platform
+
+IS_UNIX = platform.system() != "Windows"
+
+class ANSIColor(Enum):
+    RESET = '\033[0m'
+    BLACK = '\033[30m'
+    RED = '\033[31m'
+    GREEN = '\033[32m'
+    YELLOW = '\033[33m'
+    BLUE = '\033[34m'
+    PURPLE = '\033[35m'
+    CYAN = '\033[36m'
+    WHITE = '\033[37m'
+
+    # Bold colors
+    BLACK_BOLD = '\033[1;30m'
+    RED_BOLD = '\033[1;31m'
+    GREEN_BOLD = '\033[1;32m'
+    YELLOW_BOLD = '\033[1;33m'
+    BLUE_BOLD = '\033[1;34m'
+    PURPLE_BOLD = '\033[1;35m'
+    CYAN_BOLD = '\033[1;36m'
+    WHITE_BOLD = '\033[1;37m'
+
+    # Bright (light) colors
+    BRIGHT_BLACK = '\033[90m'
+    BRIGHT_RED = '\033[91m'
+    BRIGHT_GREEN = '\033[92m'  # Bright/Light Green
+    BRIGHT_YELLOW = '\033[93m'
+    BRIGHT_BLUE = '\033[94m'
+    BRIGHT_PURPLE = '\033[95m'
+    BRIGHT_CYAN = '\033[96m'
+    BRIGHT_WHITE = '\033[97m'
+
+    UNDERLINE = '\033[4m'
+    REVERSED = '\033[7m'
+
+class Logger:
+    def __init__(self, log_file_path: Optional[str] = None, verbose: Union[Literal[0, 1, 2], bool] = 2):
+        self.log_file_path: Optional[str] = log_file_path
+        if log_file_path:
+            self.log_file = open(log_file_path, "w")
+            self.log_file.close()
+        self.verbose = verbose if isinstance(verbose, int) else 2 if verbose else 0
+        assert self.verbose is not None and 0 <= self.verbose <= 2, "Please provide a valid verbose level (0, 1 or 2) default is 2"
+
+    def _log(
+            self,
+            prefix: str,
+            message: str,
+            level: int = 1,
+            color: ANSIColor = ANSIColor.BRIGHT_GREEN,
+            text_color: ANSIColor = ANSIColor.RESET,
+            one_line : bool = False
+    ):
+        """
+        Internal logging method.
+        :param prefix: The prefix for the log message
+        :param message: The message to log
+        :param level: The verbose level (0: Essential, 1: Normal, 2: Debug)
+        :param color: The color to use for the prefix
+        """
+        if level <= self.verbose:
+            if one_line and IS_UNIX:
+                print(f"{ANSIColor.RESET.value}[{color.value}{prefix}{ANSIColor.RESET.value}] {text_color.value}{message}", end="\r")
+            else:
+                print(f"{ANSIColor.RESET.value}[{color.value}{prefix}{ANSIColor.RESET.value}] {text_color.value}{message}")
+        if self.log_file_path:
+            with open(self.log_file_path, 'a', encoding='utf-8') as log_file:
+                log_file.write(f"[{prefix}] {message}" + "\n")
+
+    def log_error(self, message: str, level: int = 1):
+        self._log("Error", message, level, ANSIColor.BRIGHT_RED, ANSIColor.BRIGHT_RED)
+
+    def log_critical(self, message: str, level: int = 1):
+        self._log("Critical", message, level, ANSIColor.RED, ANSIColor.BRIGHT_RED)
+
+    def log_exception(self, exception: Exception, level: int = 1):
+        self._log("Exception", str(exception), level, ANSIColor.BRIGHT_RED, ANSIColor.BRIGHT_RED)
+
+    def log_warning(self, message: str, level: int = 1):
+        self._log("Warning", message, level, ANSIColor.BRIGHT_YELLOW, ANSIColor.BRIGHT_YELLOW)
+
+    def log_info(self, message: str, level: int = 1):
+        self._log("Info", message, level, ANSIColor.BRIGHT_BLUE)
+
+    def log_success(self, message: str, level: int = 1):
+        self._log("Success", message, level, ANSIColor.BRIGHT_GREEN)
+
+    def log_debug(self, message: str, level: int = 2):
+        self._log("Debug", message, level, ANSIColor.BRIGHT_PURPLE)
+
+    def log_verbose(self, message: str, level: int = 2):
+        self._log("Verbose", message, level, ANSIColor.BRIGHT_CYAN)
+
+    def log_custom(
+            self,
+            prefix: str,
+            message: str,
+            level: int = 1,
+            color: ANSIColor = ANSIColor.BRIGHT_GREEN,
+            text_color: ANSIColor = ANSIColor.RESET,
+            one_line : bool = False
+    ):
+        self._log(prefix, message, level, color, text_color, one_line)
diff --git a/src/OpenHosta/core/pydantic_usage.py b/src/OpenHosta/core/pydantic_usage.py
index 807e3cc..289a8d0 100644
--- a/src/OpenHosta/core/pydantic_usage.py
+++ b/src/OpenHosta/core/pydantic_usage.py
@@ -37,7 +37,7 @@ class Func(BaseModel):
         f_doc: str = Field(default="", description="Documentation of the function, e.g., 'This function returns the sum of two integers.'")
         f_call: str = Field(default="", description="Actual call of the function, e.g., 'func(1, 'hello')'")
         f_args: Dict[str, Any] = Field(default_factory=dict, description="Arguments of the function, e.g., {'a': 1, 'b': 'hello'}")
-        f_type: Tuple[List[Any], Any] = Field(default_factory=lambda: ([], None), description="Desired type of the _inputs and _outputs of the function")
+        f_type: Tuple[List[Any], Any] = Field(default_factory=lambda: ([], None), description="Desired type of the _inputs and outputs of the function")
         f_schema: Dict[str, Any] = Field(default_factory=dict, description="Dictionary describing the function's return type (in case of pydantic).")
         f_sig: Optional[inspect.Signature] = Field(default=None, description="Signature of the function")
         f_locals: Optional[Dict[str, Any]] = Field(default=None, description="Local variables within the function's scope")
diff --git a/src/OpenHosta/exec/emulate.py b/src/OpenHosta/exec/emulate.py
index f732886..e20d245 100644
--- a/src/OpenHosta/exec/emulate.py
+++ b/src/OpenHosta/exec/emulate.py
@@ -47,7 +47,6 @@ def emulate(
         **llm_args
 ) -> Any:
     x = None
-    l_ret: Any = None
 
     if _infos is None:
         x = Hosta()
@@ -59,19 +58,19 @@ def emulate(
         model = DefaultManager.get_default_model()
 
     if x:
-        x._attach(_infos.f_obj, {
+        x.attach(_infos.f_obj, { # type: ignore
             "_last_request": None,
             "_last_response": None
         })
 
     try:
         if x:
-            x._attach(_infos.f_obj, {"_last_request": {
+            x.attach(_infos.f_obj, {"_last_request": { # type: ignore
                     'sys_prompt':f"{EMULATE_PROMPT!r}\n{func_prompt}\n",
                     'user_prompt':_infos.f_call
                     }
                 }
-            )
+                     )
         response = model.simple_api_call(
             sys_prompt=f"{EMULATE_PROMPT!r}\n{func_prompt}\n",
             user_prompt=_infos.f_call,
@@ -79,7 +78,7 @@ def emulate(
         )
 
         if x:
-            x._attach(_infos.f_obj, {"_last_response": response})
+            x.attach(_infos.f_obj, {"_last_response": response}) # type: ignore
         
         l_ret = model.request_handler(response, _infos)
         if post_callback is not None:
diff --git a/src/OpenHosta/exec/example.py b/src/OpenHosta/exec/example.py
index f21f751..d8d70b2 100644
--- a/src/OpenHosta/exec/example.py
+++ b/src/OpenHosta/exec/example.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, get_args
+from typing import Any, get_args, get_origin, Literal
 
 from ..core.hosta import Hosta, ExampleType
 
@@ -15,6 +15,7 @@ def example(*args, hosta_out: Any = None, **kwargs):
     expected_return_type = x._infos.f_type[1]
 
     if not (
+        get_origin(expected_return_type) is Literal or
         isinstance(hosta_out, expected_return_type) or
         (expected_return_type == float and isinstance(hosta_out, int)) or
         hosta_out in get_args(expected_return_type) or
diff --git a/src/OpenHosta/exec/generate_data.py b/src/OpenHosta/exec/generate_data.py
new file mode 100644
index 0000000..2a3b0f1
--- /dev/null
+++ b/src/OpenHosta/exec/generate_data.py
@@ -0,0 +1,56 @@
+import inspect
+from typing import Callable, Optional, Union, Literal
+
+from .predict.dataset.dataset import HostaDataset
+from .predict.dataset.oracle import LLMSyntheticDataGenerator
+from ..core.config import Model, DefaultModel
+from ..core.hosta import Func
+from ..core.logger import Logger
+
+
+def _analyze_function(func: Callable) -> Func:
+    if not callable(func):
+        raise TypeError("The provided object is not a function or callable.")
+
+    func_obj = Func()
+    func_obj.f_name = func.__name__
+    func_obj.f_doc = func.__doc__ if func.__doc__ else ""
+    func_obj.f_sig = inspect.signature(func)
+
+    arg_types = {}
+    input_types = []
+
+    for name, param in func_obj.f_sig.parameters.items():
+        param_type = param.annotation if param.annotation != inspect.Parameter.empty else None
+        arg_types[name] = {
+            "type": param_type,
+            "default": param.default if param.default != inspect.Parameter.empty else None
+        }
+        input_types.append(param_type)
+
+    func_obj.f_args = arg_types
+    return_type = func_obj.f_sig.return_annotation if func_obj.f_sig.return_annotation != inspect.Signature.empty else None
+    func_obj.f_type = (input_types, return_type)
+    func_obj.f_mem = None
+
+    return func_obj
+
+
+def generate_data(
+        func: Callable,
+        ammount: int,
+        oracle: Optional[Model] = None,
+        verbose: Union[Literal[0, 1, 2], bool] = 2
+):
+    logger: Logger = Logger(verbose=verbose)
+    request_amounts = int(ammount / 100) if ammount > 100 else 1
+
+    logger.log_custom("Data Generation", f"Generating {ammount} examples for function {func.__name__}")
+    data = LLMSyntheticDataGenerator.generate_synthetic_data(
+        func=_analyze_function(func),
+        logger=logger,
+        request_amounts=request_amounts,
+        examples_in_req=int(ammount / request_amounts),
+        model=oracle if oracle is not None else DefaultModel().get_default_model()
+    )
+    return HostaDataset.from_list(data, logger)
diff --git a/src/OpenHosta/exec/predict/dataset/dataset.py b/src/OpenHosta/exec/predict/dataset/dataset.py
index 38b2519..eb2db01 100644
--- a/src/OpenHosta/exec/predict/dataset/dataset.py
+++ b/src/OpenHosta/exec/predict/dataset/dataset.py
@@ -1,14 +1,18 @@
-import csv
-import json
-import os
-from enum import Enum
-from typing import List, Optional, Any, Dict
+from typing import Any, Dict, List, Optional, Tuple, Union, get_origin, get_args
+from typing_extensions import Literal
 
+import os
+import csv, json
 import torch
+from torch.utils.data import DataLoader, TensorDataset, random_split
+
+from enum import Enum
 
 from .sample_type import Sample
 from ..encoder.simple_encoder import SimpleEncoder
-
+from ..predict_memory import File
+from ....core.hosta import Func
+from ....core.logger import Logger
 
 class SourceType(Enum):
     """
@@ -17,127 +21,402 @@ class SourceType(Enum):
     CSV = "csv"
     JSONL = "jsonl"
     JSON = "json"
+
 class HostaDataset:
-    def __init__(self, verbose: int = 1):
+    """
+    class for managing datasets in Hosta
+    """
+    def __init__(self, log: Logger):
         self.path: Optional[str] = None  # Path to the file
         self.data: List[Sample] = []  # List of Sample objects
-        self.dictionary: Dict[str, int] = {}  # Dictionary for mapping str to id
+        self.dictionary: Dict[str, int] = {}  # Dictionary for mapping str to id for encoding (for simple encoder -> Mini word2vec)
         self.inference: Optional[Sample] = None  # Inference data for understanding the data
-        self.verbose: int = verbose  # Verbose level for debugging
-        self._encoder: Optional[SimpleEncoder] = None  # Will store the encoder instance
+        self.log = log # Logger for logging the data
+        self.encoder: Optional[SimpleEncoder] = None  # Encoder for encoding the data
 
-    def add(self, sample: Sample):
+    ########################################################
+    ### Managing data ###
+    def add(self, data : Any, dataset : Optional[List[Sample]] = None) -> None:
         """
-        Add a Sample object to the dataset.
+        Add data to the dataset
         """
-        self.data.append(sample)
-    def convert_data(self, batch_size: int, shuffle: bool, train_set_size: float = 0.8) -> tuple:
+        if dataset is None:
+            dataset = self.data
+        
+        data_sampled = Sample(data)
+        dataset.append(data_sampled)
+        return None
+
+
+    ########################################################
+    ### Preparation of the data ###
+    def encode(self, max_tokens: int, dataset: Optional[List[Sample]] = None, inference_data: Optional[Sample] = None,
+               func : Func = None, dictionary_path : str = None ,inference : bool = False) -> List[Sample]:
         """
-        Save the dataset to a file in the specified format and convert it into dataloader for training.
+        Encode data with a token limit for str values.
         """
-        if not isinstance(self.data[0].input, torch.Tensor):
-            self.tensorify()
-        
-        inputs = torch.stack([sample.input for sample in self.data])
-        if all(sample.output is not None for sample in self.data):
-            outputs = torch.stack([sample.output for sample in self.data])
-            dataset = torch.utils.data.TensorDataset(inputs, outputs)
+        assert func is not None, "Func attribut must be provided for encoding"
+        mapping_dict: Optional[Dict[str, int]] = None
+
+
+        output_type = func.f_type[1]
+        if get_origin(output_type) is Literal:
+            mapping_dict = self._generate_mapping_dict(output_type)
+            # print("Mapping dict : ", mapping_dict)
+
+        self.get_dictionary(dictionary_path)
+        self.encoder = SimpleEncoder.init_encoder(max_tokens, self.dictionary, dictionary_path, mapping_dict, inference) #TODO: Future, we will can choose our own encoder
+
+        if inference:
+            inference_data = inference_data if inference_data is not None else self.inference
+            data_encoded = self.encoder.encode([inference_data])
+            self.inference = data_encoded[0]
         else:
-            dataset = torch.utils.data.TensorDataset(inputs)
-        
-        return self._create_dataloaders(dataset, batch_size, shuffle, train_set_size)
+            dataset = dataset if dataset is not None else self.data
+            data_encoded = self.encoder.encode(dataset)
+            self.data = data_encoded
 
-    def save_data(self, file_path: str):
+        return data_encoded
+
+
+    def decode(self, predictions: Optional[Union[List[torch.Tensor], torch.Tensor]], func_f_type: Any) -> Tuple[Any, Any]:
         """
-        Sauvegarde le dataset au format JSON
+        Decode the model predictions based on the function's return type.
         """
-        data_to_save = {
-            'data': [
-                {
-                    '_inputs': sample.input.tolist() if isinstance(sample.input, torch.Tensor) else sample.input,
-                    '_outputs': sample.output.tolist() if isinstance(sample.output, torch.Tensor) else sample.output
-                }
-                for sample in self.data
-            ]
-        }
-        with open(file_path, 'w') as f:
-            json.dump(data_to_save, f)
+        output_type = func_f_type[1]
+        output = self.encoder.decode(predictions, output_type)
+        return output, predictions
 
-    def load_data(self, file_path: str):
-        """
-        Charge un dataset depuis un fichier JSON
+
+    def tensorize(self, value: Optional[Union[List[Sample], Sample]] = None, dtype: torch.dtype = None) -> List[Sample]:
         """
-        with open(file_path, 'r') as f:
-            data_dict = json.load(f)
+        Convert data to tensors for training.
         
-        for sample_dict in data_dict['data']:
-            self.add(Sample(sample_dict))
+        Args:
+            value: Optional data to tensorize (defaults to self.data)
+            dtype: Tensor dtype (defaults to torch.float32)
         
-    def _create_dataloaders(self, dataset, batch_size: int, shuffle: bool, train_set_size: float):
-        """
-        Méthode utilitaire pour créer les dataloaders
+        Returns:
+            List of tensorized Samples
         """
-        train_size = int(train_set_size * len(dataset))
-        
-        train_dataset = torch.utils.data.Subset(dataset, range(train_size))
-        val_dataset = torch.utils.data.Subset(dataset, range(train_size, len(dataset)))
+        dtype = dtype if dtype is not None else torch.float32
+
+        if value is None:
+            value = self.data
+        elif isinstance(value, Sample):
+            value = [value]
+
+        for sample in value:
+            sample._inputs = torch.tensor(sample.input, dtype=dtype)
+
+            if sample.output is not None:
+                if isinstance(sample.output, (int, float)):
+                    sample._outputs = torch.tensor(sample.output, dtype=dtype)
+                else:
+                    sample._outputs = torch.tensor(sample.output, dtype=torch.long)
         
-        return (
-            torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle),
-            torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
-        )
+        if value[0].output is None:
+            self.inference = value[0]  # TODO: fix the error here
+        else:
+            self.data = value
 
+        return value
 
-    @staticmethod
-    def from_data(data_path: str, batch_size: int, shuffle: bool, train_set_size: float = 0.8, verbose: int = 1) -> tuple:
+    def to_dataloaders(self,  batch_size: int, data: Optional[List[Sample]] = None,
+                    train_ratio: float = 0.8, shuffle: bool = True) -> Tuple[DataLoader, DataLoader]:
         """
-        Load a dataset from a file and convert it into dataloader for training.
+        Process the data into DataLoader objects.
         """
-        dataset = HostaDataset(verbose)
-        dataset.load_data(data_path)
-        return dataset.convert_data(batch_size, shuffle, train_set_size)
+        assert 0 < train_ratio < 1, "Train ratio must be between 0 and 1"
+        assert batch_size > 0, "Batch size must be greater than 0"
 
+        data = data if data is not None else self.data
 
-    def save(self, path: str, source_type: SourceType = SourceType.CSV, elements: Optional[List[Sample]] = None):
+        # Tensorize if needed
+        if not isinstance(data[0]._inputs, torch.Tensor):
+            data = self.tensorize(data)
+
+        # Stack tensors
+        inputs = torch.stack([sample._inputs for sample in data])
+
+        if all(sample._outputs is not None for sample in data):
+            outputs = torch.stack([sample._outputs for sample in data])
+        else:
+            raise ValueError("Output data is missing in the dataset at least for one sample")
+
+        dataset = TensorDataset(inputs, outputs)
+
+        train_size = int(train_ratio * len(dataset))
+        val_size = len(dataset) - train_size
+        # print("Train size : ", train_size)
+        # print("Val size : ", val_size)
+        train_set, val_set = random_split(dataset, [train_size, val_size])
+
+        train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=shuffle)
+        val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=shuffle)
+        # print("Train loader len : ", len(train_loader))
+        # print("Val loader len : ", len(val_loader))
+        return train_loader, val_loader
+
+    def normalize_data(self, normalization_file: File, data: Optional[Union[List[Sample], Sample]] = None) -> List[Sample]:
         """
-        Save the dataset or specific elements to a file in the specified format.
-        Converts Sample objects back to dictionaries for storage.
+        Normalize the input and output data column-wise to the range [-1, 1].
 
         Args:
-            path: Path where to save the file
-            source_type: Type of file format to save (CSV, JSONL, or PICKLE)
-            elements: Optional list of Sample objects to save. If None, saves entire dataset
-        """
-        self.path = path
-        data_to_save = elements if elements is not None else self.data
-
-        # Convert Samples to dictionaries for saving
-        dict_data = []
-        for sample in data_to_save:
-            sample_dict = {}
-            for i, input_value in enumerate(sample.input):
-                sample_dict[f'input_{i}'] = input_value
+            normalization_file (File): Path to the file where the normalization stats will be saved.
+            data (Optional[Union[List[Sample], Sample]]): Data to normalize. If None, normalizes self.data.
+
+        Returns:
+            List[Sample]: Normalized samples.
+        """
+        if data is None:
+            data = self.data
+        elif isinstance(data, Sample):
+            data = [data]
+
+        if len(data) == 0:
+            raise ValueError("No data to normalize.")
+
+        num_columns = len(data[0].input)
+        min_values = [float('inf')] * num_columns
+        max_values = [float('-inf')] * num_columns
+
+        min_output = float('inf')
+        max_output = float('-inf')
+
+        for sample in data:
+            for col_idx, value in enumerate(sample.input):
+                if value < min_values[col_idx]:
+                    min_values[col_idx] = value
+                if value > max_values[col_idx]:
+                    max_values[col_idx] = value
             if sample.output is not None:
-                sample_dict['_outputs'] = sample.output
-            dict_data.append(sample_dict)
+                if sample.output < min_output:
+                    min_output = sample.output
+                if sample.output > max_output:
+                    max_output = sample.output
+
+        normalization_data = {
+            'inputs': {'min': min_values, 'max': max_values},
+            'output': {'min': min_output, 'max': max_output}
+        }
 
-        if source_type == SourceType.CSV:
-            with open(path, 'w', newline='', encoding='utf-8') as f:
-                if not dict_data:
-                    return
-                writer = csv.DictWriter(f, fieldnames=dict_data[0].keys())
-                writer.writeheader()
-                writer.writerows(dict_data)
+        for sample in data:
+            for col_idx, value in enumerate(sample.input):
+                if max_values[col_idx] == min_values[col_idx]:
+                    sample.input[col_idx] = 0
+                else:
+                    sample.input[col_idx] = 2 * (value - min_values[col_idx]) / (max_values[col_idx] - min_values[col_idx]) - 1
+            if sample.output is not None:
+                if max_output == min_output:
+                    sample.output = 0
+                else:
+                    sample.output = 2 * (sample.output - min_output) / (max_output - min_output) - 1
 
-        elif source_type == SourceType.JSONL:
-            with open(path, 'w', encoding='utf-8') as f:
-                for row in dict_data:
-                    json.dump(row, f)
-                    f.write('\n')
+        with open(normalization_file.path, 'w', encoding='utf-8') as f:
+            json.dump(normalization_data, f, indent=2, sort_keys=True)  # type: ignore
+
+        self.data = data
+        return self.data
+
+    def normalize_input_inference(self, normalization_file: File):
+        """
+        Normalize the input data column-wise to the range [-1, 1].
+
+        Args:
+            normalization_file (File): Path to the file where the normalization stats are stored.
+
+        Returns:
+            Sample: Normalized input sample.
+        """
+        if self.inference is None:
+            raise ValueError("No data to normalize.")
+
+        with open(normalization_file.path, 'r', encoding='utf-8') as f:
+            normalization_data = json.load(f)
+
+        input_min = normalization_data['inputs']['min']
+        input_max = normalization_data['inputs']['max']
+
+        for col_idx, value in enumerate(self.inference.input):
+            if input_max[col_idx] != input_min[col_idx]:
+                self.inference.input[col_idx] = 2 * (value - input_min[col_idx]) / (input_max[col_idx] - input_min[col_idx]) - 1
+            else:
+                self.inference.input[col_idx] = 0
+
+        return self.inference
+
+
+    def denormalize_output_inference(self, output: float, normalization_file: File):
+        """
+        Denormalize the output data using the normalization stats stored in the given file.
+        The file must contain the min and max for the output.
+
+        Args:
+            output (float): Output data to denormalize.
+            normalization_file (File): Path to the normalization metadata file.
+
+        Returns:
+            float: Denormalized output.
+        """
+        with open(normalization_file.path, 'r', encoding='utf-8') as f:
+            normalization_data = json.load(f)
 
+        output_min = normalization_data['output']['min']
+        output_max = normalization_data['output']['max']
+
+        if output_max != output_min:
+            output = ((output + 1) / 2) * (output_max - output_min) + output_min
         else:
-            raise ValueError(f"Unsupported source type: {source_type}")
+            output = output_min
+
+        return output
+
+    @staticmethod
+    def denormalize_output(output: Sample, normalization_file: File):
+        """
+        Denormalize the output data using the normalization stats stored in the given file.
+        The file must contain the min and max for the output.
+
+        Args:
+            output (Sample): Output data to denormalize.
+            normalization_file (File): Path to the normalization metadata file.
+
+        Returns:
+            Sample: Denormalized output sample.
+        """
+        with open(normalization_file.path, 'r', encoding='utf-8') as f:
+            normalization_data = json.load(f)
 
+        output_min = normalization_data['output']['min']
+        output_max = normalization_data['output']['max']
+
+        if output_max != output_min:
+            output.output = ((output.output + 1) / 2) * (output_max - output_min) + output_min
+        else:
+            output.output = output_min
+
+        return output
+
+    def manage_example(self):
+        pass
+
+    def examples_to_eval(self):
+        pass
+
+
+    ########################################################
+    ### Internal function process ###
+    @staticmethod
+    def _generate_mapping_dict(literal) -> dict[str, int]:
+        """
+        Generate a mapping dictionary for the output type.
+        Parameters:
+            out: A Literal type containing the possible values
+        Returns:
+            dict: A dictionary mapping each value to a unique integer (starting from 0)
+        """
+        if get_origin(literal) is not Literal:
+            raise ValueError("The literal arg must be a Literal type")
+
+        unique_values = list(get_args(literal))
+        mapping_dict = {value: idx for idx, value in enumerate(unique_values)}
+        return mapping_dict
+    
+    def get_dictionary (self, dictionary_path: str) -> None:
+        """
+        Load the tokenizer dictionary from a file
+        """
+        if not os.path.exists(dictionary_path):
+            self.dictionary = {}
+        else:
+            with open(dictionary_path, 'r', encoding='utf-8') as f:
+                loaded_dict = json.load(f)
+                self.dictionary = loaded_dict if loaded_dict else {}
+
+    def save_dictionary(self, dictionary_path: str) -> None:
+        """
+        Save the tokenizer dictionary of mapping to a file
+        """
+        # Create directory if it doesn't exist
+        os.makedirs(os.path.dirname(dictionary_path), exist_ok=True)
+
+        with open(dictionary_path, 'w', encoding='utf-8') as f:
+            json.dump(self.dictionary, f, indent=2, sort_keys=True) # type: ignore
+
+    def prepare_inference(self, inference_data: dict, max_tokens :int, func : Func, dictionary_path :str) -> None:
+            """
+            Prepare the inference data for the model.
+            """
+            self.inference = Sample(inference_data)
+            self.encode(max_tokens=max_tokens, inference_data=self.inference, func=func, dictionary_path=dictionary_path, inference=True)
+            self.tensorize(self.inference)
+
+    def save(self, path: str, source_type: SourceType = SourceType.CSV, elements: Optional[List[Sample]] = None):
+            """
+            Save the dataset or specific elements to a file in the specified format.
+            Converts Sample objects back to dictionaries for storage.
+
+            Args:
+                path: Path where to save the file
+                source_type: Type of file format to save (CSV, JSONL, or PICKLE)
+                elements: Optional list of Sample objects to save. If None, saves entire dataset
+            """
+            self.path = path
+            data_to_save = elements if elements is not None else self.data
+
+            # Convert Samples to dictionaries for saving
+            dict_data = []
+            for sample in data_to_save:
+                sample_dict = {}
+                for i, input_value in enumerate(sample.input):
+                    sample_dict[f'input_{i}'] = input_value
+                if sample.output is not None:
+                    sample_dict['outputs'] = sample.output
+                dict_data.append(sample_dict)
+
+            if source_type == SourceType.CSV:
+                with open(path, 'w', newline='', encoding='utf-8') as f:
+                    if not dict_data:
+                        return
+                    writer = csv.DictWriter(f, fieldnames=dict_data[0].keys()) # type: ignore
+                    writer.writeheader()
+                    writer.writerows(dict_data)
+
+            elif source_type == SourceType.JSONL:
+                with open(path, 'w', encoding='utf-8') as f:
+                    for row in dict_data:
+                        json.dump(row, f) # type: ignore
+                        f.write('\n')
+
+            else:
+                raise ValueError(f"Unsupported source type: {source_type}")
+            
+    def save_data(self, file_path: str):
+        """
+        Sauvegarde le dataset au format JSON
+        """
+        data_to_save = {
+            'data': [
+                {
+                    'inputs': sample.input.tolist() if isinstance(sample.input, torch.Tensor) else sample.input,
+                    'outputs': sample.output.tolist() if isinstance(sample.output, torch.Tensor) else sample.output
+                }
+                for sample in self.data
+            ]
+        }
+        with open(file_path, 'w', encoding='utf-8') as f:
+            json.dump(data_to_save, f) # type: ignore
+
+    def load_data(self, file_path: str):
+        with open(file_path, 'r', encoding='utf-8') as f:
+            data_dict = json.load(f)
+        
+        for sample_dict in data_dict['data']:
+            self.add(sample_dict)
+
+    ########################################################
+    ### Conversion ##
     def convert_files(self, path: str, source_type: Optional[SourceType] = None) -> List[Sample]:
         """
         Load dataset from a file and convert each row to a Sample object.
@@ -158,14 +437,15 @@ def convert_files(self, path: str, source_type: Optional[SourceType] = None) ->
                 source_type = SourceType.CSV
             elif path.endswith('.jsonl'):
                 source_type = SourceType.JSONL
+            elif path.endswith('.json'):
+                source_type = SourceType.JSON
             else:
                 raise ValueError(f"Please specify the source type for the file: {path}")
 
         if source_type == SourceType.CSV:
-            with open(path, 'r') as f:
+            with open(path, 'r', encoding='utf-8') as f:
                 reader = csv.DictReader(f)
                 for row in reader:
-                    # Convert string numbers to float if possible
                     processed_row = {}
                     for key, value in row.items():
                         try:
@@ -175,163 +455,127 @@ def convert_files(self, path: str, source_type: Optional[SourceType] = None) ->
                     self.data.append(Sample(processed_row))
 
         elif source_type == SourceType.JSONL:
-            with open(path, 'r') as f:
+            with open(path, 'r', encoding='utf-8') as f:
                 for line in f:
                     record = json.loads(line)
                     if not isinstance(record, dict):
                         record = {'input_0': record}
                     self.data.append(Sample(record))
 
+        elif source_type == SourceType.JSON:
+            with open(path, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+                if isinstance(data, dict):
+                    self.convert_dict(data)
+                elif isinstance(data, list):
+                    self.convert_list(data)
+                else:
+                    raise ValueError(f"Unsupported data format in JSON file: {path}")
+
         else:
             raise ValueError(f"Unsupported source type: {source_type}")
         return self.data
 
-    def convert_list(self, data: list) -> List[Sample]:
-        """
-        Create a dataset from a list.
-
-        Args:
-            data: List of dictionaries or tuples/lists representing Sample _inputs and _outputs.
-                  Each item should either be:
-                  - a dict with keys for _inputs (e.g., 'input_0', 'input_1', ...) and optional '_outputs', or
-                  - a tuple/list where the first part is _inputs(s) and the last item is _outputs (optional).
 
-        Returns:
-            HostaDataset instance
+    def convert_list(self, data: list = None) -> List[Sample]:
         """
-
+        Convert a dataset from a list.
+        """
+        data = data if data is not None else self.data
+        output_data = []
         for entry in data:
             if isinstance(entry, dict):
                 # If the entry is already a dictionary, let's assume it has the keys in the right structure
-                self.add(Sample(entry))
+                output_data.append(Sample(entry))
             elif isinstance(entry, (list, tuple)):
-                # If it's a list or tuple, we assume it's structured as (_inputs..., [_outputs])
-                inputs = list(entry[:-1])  # All but last element are _inputs
-                output = entry[-1] if len(entry) > 1 else None  # Last element could be _outputs if present
+                # If it's a list or tuple, we assume it's structured as (inputs..., [outputs])
+                inputs = list(entry[:-1])  # All but last element are inputs
+                output = entry[-1] if len(entry) > 1 else None  # Last element could be outputs if present
                 sample_dict = {f'input_{i}': input_value for i, input_value in enumerate(inputs)}
                 if output is not None:
-                    sample_dict['_outputs'] = output
-                self.add(Sample(sample_dict))
+                    sample_dict['outputs'] = output
+                output_data.append(Sample(sample_dict))
             else:
                 raise ValueError(f"Unsupported data format in list entry: {entry}")
 
-    def encode(self, max_tokens: int) -> None:
-        """
-        Encode le dataset d'entraînement et crée le dictionnaire
-        """
-        if self._encoder is None:
-            self._encoder = SimpleEncoder()
-        self.data = self._encoder.encode(self.data, max_tokens=max_tokens)
-        self.dictionary = self._encoder.dictionary
-
-    def encode_inference(self) -> None:
-        """
-        Encode les données d'inférence avec le dictionnaire existant
-        """
-        if self.dictionary is None:
-            raise ValueError("No dictionary available. Call encode() first on training data")
-        
-        self._encoder = SimpleEncoder(existing_dict=self.dictionary)
-        self.inference = self._encoder.encode([self.inference], max_tokens=10)[0]
+        self.data = output_data
+        return self.data
 
-    def tensorify(self, dtype=None) -> None:
-        """
-        Convertit le dataset d'entraînement en tenseurs
-        """
-        if dtype is None:
-            dtype = torch.float32
-            
-        for sample in self.data:
-            if not isinstance(sample.input, torch.Tensor):
-                sample._inputs = torch.tensor(sample.input, dtype=dtype)
-            
-            if sample.output is not None and not isinstance(sample.output, torch.Tensor):
-                if isinstance(sample.output, (int, float)):
-                    sample._outputs = torch.tensor(sample.output, dtype=dtype)
-                else:
-                    sample._outputs = torch.tensor(sample.output, dtype=dtype)
 
-    def tensorify_inference(self, dtype=None) -> None:
+    def convert_dict(self, data: dict) -> List[Sample]:
         """
-        Convertit les données d'inférence en tenseurs
+        Convert a dataset from a dict.
         """
-        if dtype is None:
-            dtype = torch.float32
-        
-        if not isinstance(self.inference.input, torch.Tensor):
-            self.inference._inputs = torch.tensor(self.inference.input, dtype=dtype)
+        data = data if data is not None else self.data
+        output_data = []
+        for key, value in data.items():
+            if not isinstance(value, dict):
+                value = {'outputs': value}
+            value['outputs'] = value.pop(key, None)
+            output_data.append(Sample(value))
+        self.data = output_data
+        return self.data
 
-    def prepare_inference(self, inference_data: dict) -> None:
-        """
-        Prépare les données d'inférence en les encodant et les convertissant en tenseurs
-        """
-        self.inference = Sample(inference_data)
-        self.encode_inference()
-        self.tensorify_inference()
 
+    ########################################################
+    ### Class Generator ###
     @staticmethod
-    def from_input(inference_data: dict, verbose: int = 0) -> 'HostaDataset':
+    def from_files(path: str, source_type: Optional[SourceType], log: Logger) -> 'HostaDataset':
         """
-        Crée un dataset à partir de données d'inférence
+        Load a dataset from a file.
         """
-        dataset = HostaDataset(verbose)
-        dataset.prepare_inference(inference_data)
+        dataset = HostaDataset(log)
+        dataset.convert_files(path, source_type)
         return dataset
-
-    def decode(self, predictions: List[Any], func_f_type: Any) -> List[Any]:
+    
+    @staticmethod
+    def from_list(data: list, logger: Logger) -> 'HostaDataset':
         """
-        Decode the model predictions based on the function's return type.
+        Create a dataset from a list.
         """
-        if self._encoder is None:
-            raise ValueError("Dataset must be encoded before decoding")
-        
-        # Check if func_f_type is a typing.Literal
-        # if isinstance(func_f_type, typing._GenericAlias) and get_origin() is Literal:
-
-        # if get_origin(func_f_type) is Literal:
-            # Return decoded predictions using the encoder
-            # return [self._encoder.decode_prediction(pred) for pred in predictions]
-        # else:
-        decoded_predictions = []
-        for pred in predictions:
-            pred_value = pred.detach().cpu().numpy()
-            # Convert pred_value to the expected type
-            # Handle scalar and array predictions
-            if pred_value.size == 1:
-                pred_scalar = pred_value.item()
-            else:
-                pred_scalar = pred_value
-            try:
-                converted_pred = func_f_type(pred_scalar)
-            except (TypeError, ValueError):
-                converted_pred = pred_scalar  # Return as is if conversion fails
-            decoded_predictions.append(converted_pred)
-            if func_f_type != list:
-                decoded_predictions = decoded_predictions[0]
-        return decoded_predictions
-        
+        dataset = HostaDataset(logger)
+        dataset.convert_list(data)
+        return dataset
+    
     @staticmethod
-    def from_files(path: str, source_type: Optional[SourceType], verbose: int = 1) -> 'HostaDataset':
+    def from_input(inference_data: dict, logger: Logger, max_tokens : int, func: Func, dictionary_path : str) -> 'HostaDataset':
         """
-        Load a dataset from a file.
+        Crée un dataset à partir de données d'inférence
         """
-        dataset = HostaDataset(verbose)
-        dataset.convert_files(path, source_type)
+        dataset = HostaDataset(logger)
+        dataset.prepare_inference(inference_data, max_tokens, func, dictionary_path)
         return dataset
-    
+
     @staticmethod
-    def from_list(data: list, verbose: int) -> 'HostaDataset':
+    def from_data(data_path: str, logger: Logger,) -> 'HostaDataset':
         """
-        Create a dataset from a list.
+        Load a dataset from a file and convert it into dataloader for training.
         """
-        dataset = HostaDataset(verbose)
-        dataset.convert_list(data)
+        dataset = HostaDataset(logger)
+        dataset.load_data(data_path)
+
         return dataset
+        # if config.batch_size is None:
+        #     config.batch_size = max(1, min(16384, int(0.05 * len(dataset.data))))
 
+        # return dataset.to_dataloaders(batch_size=config.batch_size, shuffle=shuffle, train_ratio=train_ratio)
 
-    @staticmethod
     def __len__(self):
         return len(self.data)
 
     def __iter__(self):
         return iter(self.data)
+    
+
+
+
+
+if __name__ == "__main__":
+    logger = Logger()
+    dataset = HostaDataset(logger)
+
+    dataset.convert_files("data.csv", SourceType.CSV)
+
+    dataset.encode(...)
+    dataset.tensorize()
+    train, val = dataset.to_dataloaders(32)
diff --git a/src/OpenHosta/exec/predict/dataset/oracle.py b/src/OpenHosta/exec/predict/dataset/oracle.py
index abccebd..5afaad5 100644
--- a/src/OpenHosta/exec/predict/dataset/oracle.py
+++ b/src/OpenHosta/exec/predict/dataset/oracle.py
@@ -1,6 +1,7 @@
 import inspect
-from typing import Optional, Dict, Any, List, Type, Union, get_args, Literal
+from typing import Optional, Dict, Any, List, Type, Union, Literal,  get_args, get_origin
 
+from ....core.logger import Logger
 from ....core.config import Model, DefaultManager
 from ....core.hosta import Func
 
@@ -13,7 +14,8 @@ class LLMSyntheticDataGenerator:
 
 
     @staticmethod
-    def _validate_row(row: str, expected_fields: List[Type]) -> Optional[List[Union[str, float]]]:
+    def _validate_row(row: str, expected_fields: List[Type], logger: Logger) -> Optional[List[Union[str, float]]]:
+        logger.log_custom("Data Generation", f"Validating row: {row}", one_line=False)
         try:
             values = row.strip().split(',')
 
@@ -40,14 +42,14 @@ def _validate_row(row: str, expected_fields: List[Type]) -> Optional[List[Union[
                     else:
                         return None
 
-                elif getattr(expected_type, '__origin__', None) is Literal:
+                elif get_origin(expected_type) is Literal:
                     valid_literals = get_args(expected_type)
-                    if value in valid_literals:
-                        result.append(value)
+                    if value.strip('"') in valid_literals:
+                        result.append(value.strip('"'))
                     else:
-                        return None  # Invalid Literal
+                        return None
 
-                elif getattr(expected_type, '__origin__', None) is Union and type(None) in get_args(expected_type):
+                elif get_origin(expected_type) is Union and type(None) in get_args(expected_type):
                     non_none_types = [t for t in get_args(expected_type) if t is not type(None)]
                     for t in non_none_types:
                         if t == int:
@@ -76,7 +78,7 @@ def _validate_row(row: str, expected_fields: List[Type]) -> Optional[List[Union[
     @staticmethod
     def _format_example(input_val: Any, output_val: Any) -> str:
         """
-        Format a single example based on _inputs/_outputs types.
+        Format a single example based on inputs/outputs types.
         """
         if isinstance(input_val, (list, tuple)):
             input_str = ','.join(map(str, input_val))
@@ -104,15 +106,17 @@ def _build_user_prompt(
             for input_val, output_val in examples.items():
                 user_prompt += f"{LLMSyntheticDataGenerator._format_example(input_val, output_val)}\n"
 
-        user_prompt += f"\n\nGenerate {examples_in_req} new DIVERSE _inputs-_outputs pairs, one per line, in CSV format"
+        user_prompt += f"\n\nGenerate {examples_in_req} new DIVERSE inputs-outputs pairs, one per line, in CSV format"
         if output_type == str:
-            user_prompt += " (remember to enclose string _outputs in quotes ex: \"_outputs\")"
+            user_prompt += " (remember to enclose string outputs in quotes ex: \"outputs\")"
         user_prompt += ":\n"
 
-        user_prompt += f"{','.join(func.f_sig.parameters.keys())},_outputs"
-        user_prompt += f"\n{','.join([str(f"\n- {a} is type {b.annotation.__name__ if b.annotation != inspect.Parameter.empty else 'Any'}")\
-                                      for a, b in func.f_sig.parameters.items()])}\n"
-        user_prompt += f"- _outputs is type {output_type.__name__}\n"
+        user_prompt += f"{','.join(func.f_sig.parameters.keys())},outputs"
+        user_prompt += "\n" + "\n".join([
+            f"- {a} is type {b.annotation.__name__ if b.annotation != inspect.Parameter.empty else 'Any'}"
+            for a, b in func.f_sig.parameters.items()
+        ]) + "\n"
+        user_prompt += f"- outputs is type {output_type.__name__}\n"
 
         return user_prompt
 
@@ -120,6 +124,7 @@ def _build_user_prompt(
     @staticmethod
     def generate_synthetic_data(
             func: Func, # The function to generate data for
+            logger: Logger, # Logger to use for logging
             request_amounts: int = 3, # Amount of requests to the model
             examples_in_req: int = 50, # Examples amount in each request
             model: Optional[Model] = None # Model to use for data generation
@@ -134,7 +139,7 @@ def generate_synthetic_data(
                 to_append = {}
                 for i, key in enumerate(input_types.keys()):
                     to_append[key] = ex_inputes[i]
-                to_append["_outputs"] = ex_output
+                to_append["outputs"] = ex_output
 
         if not model:
             model = DefaultManager.get_default_model()
@@ -160,7 +165,6 @@ def generate_synthetic_data(
         result: List[Dict] = []
         conversation_history: List = []
         attempts = 0
-        expected_fields = len(input_types) + 1
 
         conversation_history.append({
             "role": "system",
@@ -199,11 +203,11 @@ def generate_synthetic_data(
                 rows = response["choices"][0]["message"]["content"].strip().split('\n')
 
                 for row in rows:
-                    cleaned_row = LLMSyntheticDataGenerator._validate_row(row, list(input_types.values()) + [output_type])
+                    cleaned_row = LLMSyntheticDataGenerator._validate_row(row, list(input_types.values()) + [output_type], logger)
                     if cleaned_row:
                         if cleaned_row not in generated_data:
                             dictrow = dict(zip(input_types.keys(), cleaned_row[:-1]))
-                            dictrow["_outputs"] = cleaned_row[-1]
+                            dictrow["outputs"] = cleaned_row[-1]
                             result.append(dictrow)
                             generated_data.append(cleaned_row)
 
@@ -213,7 +217,7 @@ def generate_synthetic_data(
                     conversation_history = [conversation_history[0]] + conversation_history[-9:]
 
             except Exception as e:
-                print(f"Error during generation: {e} line {e.__traceback__.tb_lineno}")
+                logger.log_custom("Data Generation", f"Error during generation: {e} line {e.__traceback__.tb_lineno}")
 
             attempts += 1
 
diff --git a/src/OpenHosta/exec/predict/dataset/sample_type.py b/src/OpenHosta/exec/predict/dataset/sample_type.py
index 67277db..078dc4a 100644
--- a/src/OpenHosta/exec/predict/dataset/sample_type.py
+++ b/src/OpenHosta/exec/predict/dataset/sample_type.py
@@ -6,15 +6,14 @@
 class Sample:
     """
     A class to handle data samples for machine learning.
-    Expects a dictionary where all keys except '_outputs' are considered as _inputs.
+    Expects a dictionary where all keys except  "outputs' are considered as _inputs.
     
     Example:
         data = {
-            'feature1': [1, 2],     # Any key except '_outputs' is considered _inputs
+            'feature1': [1, 2],     # Any key except 'outputs' is considered _inputs
             'feature2': {'a': 3},    # Can contain any nested structure
             'feature3': 4,           # Can contain any primitive type
-            'feature4': BaseModel(), # Can contain Pydantic models 
-            '_outputs': 9      # Optional _outputs
+            'outputs': 9      # Optional outputs
         }
         sample = Sample(data)
     """
@@ -26,7 +25,7 @@ def __init__(self, data: dict):
         self._inputs: List[Any] = []
         self._outputs: Optional[Any] = None
         
-        output_data = data.pop('_outputs', None)
+        output_data = data.pop('outputs', None)
         if output_data is not None:
             output_flattened = self._flatten_data(output_data)
             self._outputs = output_flattened[0] if len(output_flattened) == 1 else output_flattened
@@ -58,10 +57,22 @@ def input(self) -> List[Any]:
         """Get the _inputs features"""
         return self._inputs
 
+    @input.setter
+    def input(self, value: List[Any]) -> None:
+        """Set the input features."""
+        if not isinstance(value, list):
+            raise ValueError("The 'input' property must be a list.")
+        self._inputs = value
+
     @property
     def output(self) -> Optional[Any]:
         """Get the _outputs label (None if no _outputs was provided)"""
         return self._outputs
 
+    @output.setter
+    def output(self, value: Optional[Any]) -> None:
+        """Set the output label."""
+        self._outputs = value
+
     def __repr__(self) -> str:
         return f"Sample(_inputs={self.input}, _outputs={self.output})"
diff --git a/src/OpenHosta/exec/predict/encoder/simple_encoder.py b/src/OpenHosta/exec/predict/encoder/simple_encoder.py
index a738320..a9d91ce 100644
--- a/src/OpenHosta/exec/predict/encoder/simple_encoder.py
+++ b/src/OpenHosta/exec/predict/encoder/simple_encoder.py
@@ -1,17 +1,37 @@
-from typing import List, Any, Dict, Union
+from typing import List, Any, Dict, Union, Optional
+
+import torch
+import json
+import re
 
 from .base_encoder import BaseEncoder
 from ..dataset.sample_type import Sample
 
 
-class NumericEncoder(BaseEncoder):
-    def encode(self, value: Union[int, float]) -> float:
+class IntegerEncoder(BaseEncoder):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def encode(self, value: int) -> int:
+        return int(value)
+
+    def decode(self, encoded_value: float) -> int:
+        return int(encoded_value)
+
+class FloatEncoder(BaseEncoder):
+    def __init__(self) -> None:
+        super().__init__()
+    
+    def encode(self, value: float) -> float:
         return float(value)
 
-    def decode(self, encoded_value: float) -> Union[int, float]:
-        return encoded_value
+    def decode(self, encoded_value: float) -> float:
+        return float(encoded_value)
 
 class BooleanEncoder(BaseEncoder):
+    def __init__(self) -> None:
+        super().__init__()
+
     def encode(self, value: bool) -> int:
         return int(value)
 
@@ -19,108 +39,126 @@ def decode(self, encoded_value: int) -> bool:
         return bool(encoded_value)
 
 class StringEncoder(BaseEncoder):
-    def __init__(self, existing_dict: Dict[str, int] = None):
-        """
-        Initialize with optional existing dictionary.
-        If existing_dict is provided, we're in inference mode.
-        """
-        self.inference_mode = existing_dict is not None
-        self.word_to_id = {'<UNK>': 0} if existing_dict is None else existing_dict
-        self.id_to_word = {v: k for k, v in self.word_to_id.items()}
-        self.next_id = max(self.word_to_id.values()) + 1 if self.word_to_id else 1
-        self.max_tokens = None
-
-    def set_max_tokens(self, max_tokens: int):
-        """Set maximum length for encoded sequences"""
+    def __init__(self, dictionary : Dict[str, int], max_tokens : int, inference : bool) -> None:
+        super().__init__()
+        self.dictionary = dictionary
         self.max_tokens = max_tokens
+        self.inference = inference
+        self.next_id = max(self.dictionary.values()) + 1 if self.dictionary else 1
 
-    def encode(self, value: str) -> List[int]:
-        """
-        Encode a string into a list of integers.
-        For classification (_outputs), returns a single integer.
-        For _inputs features, returns a list of integers of length max_tokens.
-        """
-        if self.max_tokens is None:
-            raise ValueError("max_tokens must be set before encoding")
+    @staticmethod
+    def string_to_list(input_string):
+        tokens = re.findall(r'[a-zA-Zà-ÿ]+|[^a-zA-Zà-ÿ\s]', input_string.lower())
+        return tokens
 
-        words = str(value).lower().strip().split()
+    def encode(self, data: str) -> list[Union[str, int]]:
+        words = self.string_to_list(data)
         encoded = []
 
         for word in words:
-            if not self.inference_mode and word not in self.word_to_id:
-                self.word_to_id[word] = self.next_id
-                self.id_to_word[self.next_id] = word
-                self.next_id += 1
-            encoded.append(self.word_to_id.get(word, 0))
+            if self.inference:
+                encoded.append(self.dictionary.get(word, 0))
+            else:
+                if word not in self.dictionary:
+                    self.dictionary[word] = self.next_id
+                    encoded.append(self.next_id)
+                    self.next_id += 1
+                else:
+                    encoded.append(self.dictionary[word])
 
         if len(encoded) > self.max_tokens:
             return encoded[:self.max_tokens]
+        
         return encoded + [0] * (self.max_tokens - len(encoded))
+    
+    def decode(self, encoded_value: int) -> str:
+        return self.dictionary.get(encoded_value, '<UNK>')
+    
+    @property
+    def get_dictionnary(self) -> Dict[int, str]:
+        return self.dictionary
+
+class MappingEncoder(BaseEncoder):
+    def __init__(self, mapping_dict: Dict[int, Any]) -> None:
+        self.mapping_dict = mapping_dict
 
-    def decode(self, encoded_value: Union[int, List[int]]) -> str:
-        """
-        Decode either a single integer (classification) or list of integers (features)
-        """
-        if isinstance(encoded_value, (int, float)):
-            return self.id_to_word.get(int(encoded_value), '<UNK>')
+    def encode(self, value: Any) -> int:
+        return self.mapping_dict[value]
 
-        words = []
-        for idx in encoded_value:
-            if idx != 0:  # Skip padding
-                words.append(self.id_to_word.get(idx, '<UNK>'))
-        return ' '.join(words)
+    def decode(self, encoded_value: int) -> Any:
+        for key, value in self.mapping_dict.items():
+            if value == encoded_value:
+                return key
+        raise ValueError(f"Unknown value: {encoded_value}")
 
 class SimpleEncoder:
-    def __init__(self, existing_dict: Dict[str, int] = None):
-        self.string_encoder = StringEncoder(existing_dict)
-        self.feature_types = {}
+    def __init__(self, max_tokens: int, dictionary: Dict[str, int], dictionary_path : str, mapping_dict: Dict[str, int], inference : bool) -> None:
         self.encoders = {
-            str: self.string_encoder,
-            int: NumericEncoder(),
-            float: NumericEncoder(),
+            str: StringEncoder(dictionary, max_tokens, inference),
+            int: IntegerEncoder(),
+            float: FloatEncoder(),
             bool: BooleanEncoder()
         }
+        self.mapping_dict = mapping_dict
+        self.dictionary_path = dictionary_path
+
+    @staticmethod
+    def init_encoder(max_tokens: int, dictionary: Dict[str, int], dictionary_path : str, mapping_dict: Dict[str, int], inference : bool) -> 'SimpleEncoder':
+
+        encoder = SimpleEncoder(max_tokens, dictionary, dictionary_path ,mapping_dict, inference)
+        return encoder
+    
+    def save_dictionary(self, dictionary: Dict[int, str]) -> None:
+        with open(self.dictionary_path, 'w', encoding='utf-8') as f:
+            json.dump(dictionary, f, indent=2, sort_keys=True) # type: ignore
+
+
+    def encode(self, samples: List[Sample]) -> List[Sample]:
 
-    def encode(self, samples: List[Sample], max_tokens: int) -> List[Sample]:
-        self.string_encoder.set_max_tokens(max_tokens)
-        
         encoded_samples = []
         for sample in samples:
             encoded_input = []
             for idx, value in enumerate(sample.input):
-                encoder = self.encoders[type(value)]
-                self.feature_types[idx] = type(value)
-                encoded_value = encoder.encode(value)
-                if isinstance(encoded_value, list):
-                    encoded_input.extend(encoded_value)
+                encoder_in = self.encoders[type(value)]
+                encoded_value_in = encoder_in.encode(value)
+                if isinstance(encoded_value_in, list):
+                    encoded_input.extend(encoded_value_in)
                 else:
-                    encoded_input.append(encoded_value)
+                    encoded_input.append(encoded_value_in)
 
             encoded_output = None
             if sample.output is not None:
-                if isinstance(sample.output, str):
-                    print("\033[93mWarning: Multiple string _outputs not supported, only using the first word will be used for _outputs\033[0m")
-                output_idx = len(sample.input)
-                encoder = self.encoders[type(sample.output)]
-                self.feature_types[output_idx] = type(sample.output)
-                encoded_output = encoder.encode(sample.output)
-                # Like multiple str _outputs not supported only use the first str _outputs
-                if isinstance(encoded_output, list):
-                    encoded_output = encoded_output[0]
-
+                if self.mapping_dict is None:
+                    encoder_out = self.encoders[type(sample.output)]
+                    encoded_value_out = encoder_out.encode(sample.output)
+                else:
+                    encoder_out = MappingEncoder(self.mapping_dict)
+                    encoded_value_out = encoder_out.encode(sample.output)
+                encoded_output = encoded_value_out
+            
             encoded_samples.append(Sample({
-                '_inputs': encoded_input,
-                '_outputs': encoded_output
+                'inputs': encoded_input,
+                'outputs': encoded_output
             }))
 
+        self.save_dictionary(self.encoders[str].get_dictionnary)
         return encoded_samples
+    
+    def decode(self, predictions: Union[Optional[List[torch.Tensor]], torch.Tensor], output_type: Any) -> Any:
 
-    def decode_prediction(self, prediction: Any, position: int) -> Any:
-        if position not in self.feature_types:
-            raise ValueError(f"Unknown feature position: {position}")
+        predictions = predictions.cpu().detach().numpy()
+        if predictions is None:
+            return None
+
+        if self.mapping_dict is None:
+            predictions = predictions[0]
+            if output_type not in self.encoders:
+                raise ValueError(f"Unknown output type in decoder: {output_type}")
+            encoder_out = self.encoders[output_type]
+            return encoder_out.decode(predictions)
+        else:
+            value = predictions.argmax()
+            encoder_out = MappingEncoder(self.mapping_dict)
+            return encoder_out.decode(value)
 
-        return self.encoders[self.feature_types[position]].decode(prediction)
 
-    @property
-    def dictionary(self) -> Dict[str, int]:
-        return self.string_encoder.word_to_id
diff --git a/src/OpenHosta/exec/predict/model/builtins/algo_architecture.py b/src/OpenHosta/exec/predict/model/builtins/algo_architecture.py
new file mode 100644
index 0000000..1f80938
--- /dev/null
+++ b/src/OpenHosta/exec/predict/model/builtins/algo_architecture.py
@@ -0,0 +1,87 @@
+import math
+from typing import List
+
+def _get_nb_layer(complexity: int):
+    """
+    Get the number of layers for a neural network based on its complexity.
+    L_total refers to the total number of layers.
+    L_up refers to the number of layers in the upward path.
+    L_down refers to the number of layers in the downward path.
+    """
+    l_total = 2 * complexity + 1
+    l_up = (l_total - 1) // 2
+    l_down = l_total - l_up - 1
+    return l_up, l_down
+
+
+def _round_to_power_of_two(value: float) -> int:
+    """
+    Round a number to the nearest power of 2.
+    """
+    return 2 ** int(round(math.log2(value)))
+
+def _get_layers_size(n_in, n_out, n_peak, l_up, l_down, growth_rate, max_layer_size) -> List[int]:
+    """
+    this function calculates the size of each layer in a neural network.
+    """
+    layers_size = [n_in]
+    current_size = n_in
+
+    # Upward Phase
+    for i in range(l_up):
+        next_size = current_size * growth_rate
+        next_size = min(next_size, n_peak, max_layer_size)
+        next_size = _round_to_power_of_two(next_size)
+        layers_size.append(int(next_size))
+        current_size = next_size
+
+    # Downward Phase
+    for i in range(l_down):
+        next_size = current_size / growth_rate
+        next_size = max(next_size, n_out)
+        next_size = _round_to_power_of_two(next_size)
+        layers_size.append(int(next_size))
+        current_size = next_size
+
+    # Ensure that the last layer size is n_out
+    layers_size[-1] = n_out
+
+    return layers_size
+
+
+def get_algo_architecture(input_size: int, output_size: int, complexity: int, growth_rate: float, max_layer_coefficient: int) -> List[int]:
+    """
+    this function generates the architecture of a neural network.
+    the complexity defines the number of layers in the network.
+    growth_rate defines the rate at which the number of neurons grows.
+    max_layer_coefficient defines the maximum size of the maximum layer.
+    """
+    n_in = input_size
+    n_out = output_size
+    n_max = max(n_in, n_out)
+
+    max_layer_size = max_layer_coefficient * n_max
+
+    l_up, l_down = _get_nb_layer(complexity)
+
+    n_peak = n_in * (growth_rate ** l_up)
+    n_peak = min(n_peak, max_layer_size)
+    n_peak = _round_to_power_of_two(n_peak)  # Assurer que n_peak est une puissance de 2
+
+    layers_size = _get_layers_size(n_in, n_out, n_peak, l_up, l_down, growth_rate, max_layer_size)
+
+    return layers_size
+
+
+# test the algorithm
+########################################################
+
+# if __name__ == "__main__":
+#     input_size = 8
+#     output_size = 2
+#     complexity = 5
+#     growth_rate = 1.5
+#     max_layer_coefficient = 150
+
+#     architecture = get_algo_architecture(input_size, output_size, complexity, growth_rate, max_layer_coefficient)
+#     print("Generated architecture:", architecture)
\ No newline at end of file
diff --git a/src/OpenHosta/exec/predict/model/builtins/classification.py b/src/OpenHosta/exec/predict/model/builtins/classification.py
index d08aabd..9482b8e 100644
--- a/src/OpenHosta/exec/predict/model/builtins/classification.py
+++ b/src/OpenHosta/exec/predict/model/builtins/classification.py
@@ -1,146 +1,156 @@
-### NOT IMPLEMENTED YET ###
-
 from typing import Optional
 
 import torch
 from torch import nn
 from torch import optim
+from torch.optim.lr_scheduler import StepLR
 
+from .algo_architecture import get_algo_architecture
 from ..hosta_model import HostaModel
 from ..neural_network import NeuralNetwork
+from ..neural_network_types import ArchitectureType
+from ....predict.predict_config import PredictConfig
 from .....utils.torch_nn_utils import custom_optimizer_to_pytorch, custom_loss_to_pytorch, custom_layer_to_pytorch
-
+from .....core.logger import Logger, ANSIColor
 
 class Classification(HostaModel):
-    def __init__(self, neural_network: Optional[NeuralNetwork], input_size: int, output_size: int, complexity: int, num_classes: int, device: Optional[str] = None):
+    def __init__(
+            self,
+            neural_network: Optional[NeuralNetwork],
+            input_size: int,
+            output_size: int,
+            config: PredictConfig,
+            logger: Logger,
+            device: Optional[str] = None
+    ):
         super().__init__(device)
 
-        self.complexity = complexity
-        self.num_classes = num_classes
-        self.verbose = True
-        self.layers = []
+        self.complexity = config.complexity
+        self.logger = logger
+        self.device = device
+        self.growth_rate = config.growth_rate
+        self.max_layer_coefficent = config.coef_layers
+        self.architecture_type = ArchitectureType.CLASSIFICATION
+
         if neural_network is None or neural_network.layers is None or len(neural_network.layers) == 0:
-            transition_value = int(((input_size * output_size) / 2) * self.complexity)
-
-            input_layer = int(input_size * (2 * self.complexity))
-            if input_size > output_size:
-                hidden_layer_1 = int(transition_value / output_size)
-            else:
-                hidden_layer_1 = transition_value
-
-            # Define simple fully connected architecture
-            self.layers.append(nn.Linear(input_size, input_layer))
-            self.layers.append(nn.ReLU())  # Apply ReLU after first layer
-            self.layers.append(nn.Linear(input_layer, hidden_layer_1))
-            self.layers.append(nn.ReLU())  # Apply ReLU after second layer
-            self.layers.append(nn.Linear(hidden_layer_1, output_size))
-        else:
-            # Use custom user-defined layers from neural network definition if available
-            self.layers = [custom_layer_to_pytorch(layer) for layer in neural_network.layers]
+            layer_size : list = get_algo_architecture(input_size, output_size, self.complexity, self.growth_rate, self.max_layer_coefficent)
+
+            layers = []
+            for i in range(len(layer_size) - 1):
+                in_features = layer_size[i]
+                out_features = layer_size[i + 1]
 
-        for i, layer in enumerate(self.layers):
-            setattr(self, f'fc{i + 1}', layer)
+                linear_layer = nn.Linear(in_features, out_features)
+                layers.append(linear_layer)
+
+                if i < len(layer_size) - 2:
+                    activation = nn.ReLU()
+                    layers.append(activation)
+            self.model = nn.Sequential(*layers)
+        else:
+            layers = [custom_layer_to_pytorch(layer) for layer in neural_network.layers]
+            self.model = nn.Sequential(*layers)
 
-        # Set the loss function for classification
         if neural_network is None or neural_network.loss_function is None:
-            if num_classes == 2:
-                self.loss = nn.BCEWithLogitsLoss()  # For binary classification
-            else:
-                self.loss = nn.CrossEntropyLoss()  # For multi-class classification
+            self.loss = nn.CrossEntropyLoss()
         else:
             self.loss = custom_loss_to_pytorch(neural_network.loss_function)
 
-        # Set the optimizer
         if neural_network is None or neural_network.optimizer is None:
-            self.optimizer = optim.Adam(self.parameters(), lr=0.001)
+            self.optimizer = optim.AdamW(self.parameters(), lr=0.001)
         else:
-            self.optimizer = custom_optimizer_to_pytorch(neural_network.optimizer, self, lr=0.001)
+            self.optimizer = custom_optimizer_to_pytorch(neural_network.optimizer, self, lr=0.001) # TODO: Add learning rate parameter
+
+        self.scheduler = StepLR(self.optimizer, step_size=10, gamma=0.1)
 
-        # Move model to the selected device (CPU or GPU)
         self.to(self.device)
 
 
     def trainer(self, train_set, epochs):
+        """
+        Train the model on the training set for a classification task
+        """
         self.train()
 
         for epoch in range(epochs):
+
             running_loss = 0.0
-            correct = 0
-            total = 0
+            correct_predictions = 0
+            total_samples = 0
+    
             for inputs, labels in train_set:
-                inputs, labels = inputs.to(self.device), labels.to(self.device)
+                inputs = inputs.to(self.device)
+                labels = labels.to(self.device).long() # Ensure it's long for cross entropy loss ! 
+                batch_size = inputs.size(0)
 
-                # Zero parameter gradients
                 self.optimizer.zero_grad()
-
-                # Forward pass
-                outputs = self(inputs)
-
-                if self.num_classes == 2:
-                    preds = (torch.sigmoid(outputs) > 0.5).float()
-                else:
-                    preds = torch.argmax(outputs, dim=1)
-
-                # Compute Loss
+                outputs = self.model(inputs)
                 loss = self.loss(outputs, labels)
+
                 loss.backward()
                 self.optimizer.step()
 
                 running_loss += loss.item()
-
-                # Calculate accuracy
-                if self.num_classes == 2:
-                    correct += (preds == labels).sum().item()
-                else:
-                    correct += (preds == labels.argmax(dim=1)).sum().item()
-
-                total += labels.size(0)
-
-            accuracy = correct / total
-            if self.verbose:
-                print(f"Epoch {epoch + 1}/{epochs}, Loss: {running_loss / len(train_set):.4f}, Accuracy: {accuracy * 100:.2f}%")
-
+                predicted_classes = torch.argmax(outputs, dim=1)
+
+                correct_predictions += (predicted_classes == labels).sum().item()
+                total_samples += batch_size
+            
+            self.scheduler.step()
+            current_lr = self.optimizer.param_groups[0]['lr']
+
+            epoch_loss = running_loss / len(train_set)
+            epoch_accuracy = (correct_predictions / total_samples) * 100
+            if epoch == epochs - 1:
+                self.logger.log_custom("Epoch", f"{epoch + 1}/{epochs}, Loss: {epoch_loss:.4f}, Accuracy: {epoch_accuracy:.2f}%, LR: {current_lr:.6f}", color=ANSIColor.CYAN, level=1, one_line=False)
+            else :    
+                self.logger.log_custom("Epoch", f"{epoch + 1}/{epochs}, Loss: {epoch_loss:.4f}, Accuracy: {epoch_accuracy:.2f}%, LR: {current_lr:.6f}", color=ANSIColor.CYAN, level=1, one_line=True)
+    
     def validate(self, validation_set):
-        """Validate the model's performance"""
-        self.eval()  # Set model to evaluation mode
+        """
+        Validate the model on the validation set for a classification task.
+        """
+        self.eval()
         validation_loss = 0.0
-        correct = 0
-        total = 0
+        correct_predictions = 0
+        total_samples = 0
+
         with torch.no_grad():
             for inputs, labels in validation_set:
-                inputs, labels = inputs.to(self.device), labels.to(self.device)
-                outputs = self(inputs)
+                inputs = inputs.to(self.device)
+                labels = labels.to(self.device).long() # Ensure that is long for loss
+                batch_size = inputs.size(0)
 
+                outputs = self.model(inputs)
                 loss = self.loss(outputs, labels)
-                validation_loss += loss.item()
+                validation_loss += loss.item() * batch_size
 
-                # For Classification Metrics (like binary or multi-class accuracy)
-                if self.num_classes == 2:
-                    # Binary classification: Apply sigmoid and threshold at 0.5
-                    preds = (torch.sigmoid(outputs) > 0.5).float()
-                    correct += (preds == labels).sum().item()
-                else:
-                    # Multi-class classification: Use argmax to get class labels
-                    preds = torch.softmax(outputs, dim=1)
-                    correct += (preds == labels.argmax(dim=1)).sum().item()
+                predicted_classes = torch.argmax(outputs, dim=1)
 
-                total += labels.size(0)
+                correct_predictions += (predicted_classes == labels).sum().item()
+                total_samples += labels.size(0)
 
-        avg_val_loss = validation_loss / len(validation_set)
-        accuracy = correct / total
-        print(f"Validation Loss: {avg_val_loss:.4f}, Accuracy: {accuracy * 100:.2f}%")
+            avg_val_loss = validation_loss / total_samples
+            accuracy = (correct_predictions / total_samples) * 100
 
-        return avg_val_loss, accuracy
+            self.logger.log_custom("Validation", f"Loss: {avg_val_loss:.4f}, Accuracy: {accuracy:.2f}%", color=ANSIColor.CYAN, level=1)
 
+            return # Don't need to return something for now
 
     def inference(self, x):
-        """Make prediction on a _inputs inference the model"""
+        """
+        Make prediction on inputs using the model.
+        """
         self.eval()
         with torch.no_grad():
             x = x.to(self.device)
-            outputs = self(x)
-            if self.num_classes == 2:
-                prediction = (torch.sigmoid(outputs) > 0.5).float()
-            else:
-                prediction = torch.softmax(outputs, dim=1)
-            return prediction.cpu()
+            
+            # Add batch dimension if needed
+            if x.dim() == 1:
+                x = x.unsqueeze(0)
+                    
+            outputs = self.model(x)
+
+            probabilities = torch.softmax(outputs, dim=1)
+
+            return probabilities.cpu()
\ No newline at end of file
diff --git a/src/OpenHosta/exec/predict/model/builtins/linear_regression.py b/src/OpenHosta/exec/predict/model/builtins/linear_regression.py
index 9838899..135cd43 100644
--- a/src/OpenHosta/exec/predict/model/builtins/linear_regression.py
+++ b/src/OpenHosta/exec/predict/model/builtins/linear_regression.py
@@ -3,95 +3,155 @@
 import torch
 from torch import nn
 from torch import optim
+from torch.optim.lr_scheduler import StepLR
 
+from .algo_architecture import get_algo_architecture
 from ..hosta_model import HostaModel
 from ..neural_network import NeuralNetwork
+from ..neural_network_types import ArchitectureType
+from ....predict.predict_config import PredictConfig
 from .....utils.torch_nn_utils import custom_optimizer_to_pytorch, custom_loss_to_pytorch, custom_layer_to_pytorch
+from .....core.logger import Logger, ANSIColor
 
 
 class LinearRegression(HostaModel):
-    def __init__(self, neural_network: Optional[NeuralNetwork], input_size: int, output_size: int, complexity: int, device: Optional[str] = None):
+    def __init__(
+            self,
+            neural_network: Optional[NeuralNetwork],
+            input_size: int,
+            output_size: int,
+            config: PredictConfig,
+            logger: Logger,
+            device: Optional[str] = None
+    ):
         super().__init__(device)
 
-        self.complexity = complexity
+        self.complexity = config.complexity
+        self.logger = logger
+        self.device = device
+        self.growth_rate = config.growth_rate
+        self.max_layer_coefficent = config.coef_layers
+        self.architecture_type = ArchitectureType.LINEAR_REGRESSION
 
-        self.layers = []
         if neural_network is None or neural_network.layers is None or len(neural_network.layers) == 0:
-            transition_value = int(((input_size * output_size) / 2) * self.complexity)
-
-            input_layer = int(input_size * (2 * self.complexity))
-            if input_size > output_size:
-                hidden_layer_1 = int(transition_value / output_size)
-            else:
-                hidden_layer_1 = transition_value
-
-            self.layers.append(nn.Linear(input_size, input_layer))
-            self.layers.append(nn.ReLU())
-            self.layers.append(nn.Linear(input_layer, hidden_layer_1))
-            self.layers.append(nn.ReLU())
-            self.layers.append(nn.Linear(hidden_layer_1, output_size))
-        else:
-            self.layers = [custom_layer_to_pytorch(layer) for layer in neural_network.layers]
+            layer_size : list = get_algo_architecture(input_size, output_size, self.complexity, self.growth_rate, self.max_layer_coefficent)
+
+            layers = []
+            for i in range(len(layer_size) - 1):
+                in_features = layer_size[i]
+                out_features = layer_size[i + 1]
+
+                linear_layer = nn.Linear(in_features, out_features)
+                layers.append(linear_layer)
 
-        for i, layer in enumerate(self.layers):
-            setattr(self, f'fc{i + 1}', layer)
+                if i < len(layer_size) - 2:
+                    activation = nn.ReLU()
+                    layers.append(activation)
+            self.model = nn.Sequential(*layers)
+        else:
+            layers = [custom_layer_to_pytorch(layer) for layer in neural_network.layers]
+            self.model = nn.Sequential(*layers)
+        
 
-        # Set the loss function
         if neural_network is None or neural_network.loss_function is None:
-            self.loss = nn.MSELoss()
+            self.loss = nn.SmoothL1Loss()
         else:
             self.loss = custom_loss_to_pytorch(neural_network.loss_function)
 
-        # Set the optimizer
         if neural_network is None or neural_network.optimizer is None:
-            self.optimizer = optim.Adam(self.parameters(), lr=0.001)
+            self.optimizer = optim.AdamW(self.parameters(), lr=0.01, weight_decay=1e-2)
         else:
             self.optimizer = custom_optimizer_to_pytorch(neural_network.optimizer, self, lr=0.001)
 
-        # Move model to the selected device (CPU or GPU)
+        self.scheduler = StepLR(self.optimizer, step_size=10, gamma=0.1)
+
         self.to(self.device)
 
-    def trainer(self, train_set, epochs, verbose=False):
+    def trainer(self, train_set, epochs):
+        """
+        Training loop for regression models.
+        Includes accuracy calculation with a tolerance (epsilon).
+        """
+        epsilon = 0.1
         self.train()
 
         for epoch in range(epochs):
             running_loss = 0.0
+            correct_predictions = 0
+            total_samples = 0
+            
             for inputs, labels in train_set:
-                # Move _inputs and labels to the right device
-                inputs, labels = inputs.to(self.device), labels.to(self.device)
+                inputs = inputs.to(self.device)
+                labels = labels.to(self.device).float().view(-1, 1)  # Ensure proper shape and type
+                batch_size = labels.size(0)
 
-                # Zero the parameter gradients
                 self.optimizer.zero_grad()
-
-                # Forward pass
-                outputs = self(inputs)
-                
+                outputs = self.model(inputs)
                 loss = self.loss(outputs, labels)
 
-                # Backward pass and update
                 loss.backward()
                 self.optimizer.step()
 
                 running_loss += loss.item()
-            # if verbose:
-            print(f"Epoch {epoch + 1}/{epochs}, Loss: {running_loss / len(train_set)}")
+
+                correct_predictions += ((outputs - labels).abs() < epsilon).sum().item()
+                total_samples += batch_size
+            
+            self.scheduler.step()
+            current_lr = self.optimizer.param_groups[0]['lr']
+
+            epoch_loss = running_loss / len(train_set)
+            epoch_accuracy = (correct_predictions / total_samples) * 100
+            if epoch == epochs - 1:
+                self.logger.log_custom("Epoch", f"{epoch + 1}/{epochs}, Loss: {epoch_loss:.4f}, Accuracy: {epoch_accuracy:.2f}%, LR: {current_lr:.6f}", color=ANSIColor.CYAN, level=1, one_line=False)
+            else :    
+                self.logger.log_custom("Epoch", f"{epoch + 1}/{epochs}, Loss: {epoch_loss:.4f}, Accuracy: {epoch_accuracy:.2f}%, LR: {current_lr:.6f}", color=ANSIColor.CYAN, level=1, one_line=True)
+
+
 
     def validate(self, validation_set):
-        """Validate the model on a given validation set."""
-        self.eval()  # Set model to eval mode (disable dropout, etc.)
+        """
+        Validate the model on a given validation set for regression models.
+        Includes accuracy calculation with a tolerance (epsilon).
+        """
+        self.eval()
         validation_loss = 0.0
-        with torch.no_grad():  # No need to track gradients during validation
+        correct_predictions = 0
+        total_samples = 0
+        epsilon = 0.1
+
+        with torch.no_grad():
             for inputs, labels in validation_set:
-                inputs, labels = inputs.to(self.device), labels.to(self.device)
-                outputs = self(inputs)
+                inputs = inputs.to(self.device)
+                labels = labels.to(self.device).float().view(-1, 1)
+                batch_size = labels.size(0)
+
+                outputs = self.model(inputs)
                 loss = self.loss(outputs, labels)
-                validation_loss += loss.item()
-        return validation_loss / len(validation_set)
+                validation_loss += loss.item() * batch_size
+
+                correct_predictions += ((outputs - labels).abs() < epsilon).sum().item()
+                total_samples += batch_size
+
+            avg_val_loss = validation_loss / len(validation_set)
+            accuracy = (correct_predictions / total_samples) * 100
+
+            self.logger.log_custom("Validation", f"Loss: {avg_val_loss:.4f}, Accuracy: {accuracy:.2f}%", color=ANSIColor.CYAN, level=1, one_line=False)
+
+        return #Nothing to return for now
+
 
     def inference(self, x):
-        """Make predictions for the given test set."""
+        """
+        Make predictions on a given input for linear regression task.
+        """
         self.eval()
         with torch.no_grad():
             x = x.to(self.device)
-            outputs = self(x)
+
+            if x.dim() == 1:
+                x = x.unsqueeze(0)
+
+            outputs = self.model(x)
+
             return outputs
diff --git a/src/OpenHosta/exec/predict/model/hosta_model.py b/src/OpenHosta/exec/predict/model/hosta_model.py
index 08df057..8a3338c 100644
--- a/src/OpenHosta/exec/predict/model/hosta_model.py
+++ b/src/OpenHosta/exec/predict/model/hosta_model.py
@@ -3,21 +3,24 @@
 
 import torch
 from torch import nn
+from .neural_network_types import ArchitectureType
 
 
 class HostaModel(ABC, nn.Module):
     def __init__(self, device: Optional[str]):
         self.device = device if device is not None else ('cuda' if torch.cuda.is_available() else 'cpu')
         self.layers = []
+        self.architecture_type: Optional[ArchitectureType] = None
         super().__init__()
 
     def trainer(self, train_set, epochs):
         pass
 
     def forward(self, x):
-        for layer in self.layers:
-            x = layer(x)
-        return x
+        # for layer in self.layers:
+        #     x = layer(x)
+        # return x
+        pass
 
     def validate(self, validation_set):
         pass
@@ -30,4 +33,4 @@ def init_weights(self, path: str):
         self.eval()
 
     def save_weights(self, path: str):
-        torch.save(self.state_dict(), path, )
+        torch.save(self.state_dict(), path)
diff --git a/src/OpenHosta/exec/predict/model/model_provider.py b/src/OpenHosta/exec/predict/model/model_provider.py
index dc221c2..127580b 100644
--- a/src/OpenHosta/exec/predict/model/model_provider.py
+++ b/src/OpenHosta/exec/predict/model/model_provider.py
@@ -7,28 +7,52 @@
 from .neural_network_types import ArchitectureType
 from ..predict_config import PredictConfig
 from ....core.hosta import Func
+from ....core.logger import Logger, ANSIColor
 from ....utils.torch_nn_utils import type_size
 
 
 class HostaModelProvider:
     @staticmethod
-    def from_hosta_func(func: Func, config: Optional[PredictConfig], architecture: Optional[NeuralNetwork], path: str, verbose: int) -> Optional[HostaModel]:
+    def from_hosta_func(func: Func, config: Optional[PredictConfig], architecture: Optional[NeuralNetwork], path: str, logger: Logger) -> Optional[HostaModel]:
         input_size = 0
         for arg in func.f_type[0]:
-            input_size += type_size(arg, config.max_tokens)
+            if get_origin(arg) is Literal:
+                input_size += 1
+            else:
+                input_size += type_size(arg, config.max_tokens)
+
         output_size = type_size(func.f_type[1], config.max_tokens)
+        logger.log_debug(f"Model with input size {input_size} and output size {output_size}", level=2)
+
         hosta_model: Optional[HostaModel] = None
-        if config is not None and config.model_type is not None:
-            if config.model_type == ArchitectureType.LINEAR_REGRESSION:
-                hosta_model = LinearRegression(architecture, input_size, output_size, config.complexity)
-            elif config.model_type == ArchitectureType.CLASSIFICATION:
-                hosta_model = Classification(architecture, input_size, output_size, config.complexity, 1)
+        model_type = determine_model_type(func)
+
+        if model_type == ArchitectureType.LINEAR_REGRESSION:
+            hosta_model = LinearRegression(architecture, input_size, output_size, config, logger)
+
+        elif model_type == ArchitectureType.CLASSIFICATION:
+            hosta_model = Classification(architecture, input_size, output_size, config, logger)
         else:
-            if get_origin(func.f_type[1]) == Literal:
-                hosta_model = Classification(architecture, input_size, output_size, 4, 1)
-            else:
-                hosta_model = LinearRegression(architecture, input_size, output_size, 4)
+            raise ValueError(f"Model type {model_type} not supported")
+
+        logger.log_custom("Model", f"Type : {type(hosta_model).__name__}", color=ANSIColor.BLUE_BOLD)
+
+
+        if architecture is None:
+            save_architecture(hosta_model, path, logger)
 
-        with open(path, 'w') as file:
-            file.write(NeuralNetwork.from_torch_nn(hosta_model).to_json())
         return hosta_model
+
+
+def determine_model_type(func: Func) -> ArchitectureType:
+    if get_origin(func.f_type[1]) is Literal:
+        return ArchitectureType.CLASSIFICATION
+    else:
+        return ArchitectureType.LINEAR_REGRESSION
+
+
+def save_architecture(hosta_model: HostaModel, path: str, logger : Logger):
+    architecture = NeuralNetwork.from_torch_nn(hosta_model)
+    with open(path, 'w', encoding='utf-8') as file:
+        file.write(architecture.to_json())
+    logger.log_custom("Architecture", f"saved to {path}", color=ANSIColor.BRIGHT_GREEN, level=2)
\ No newline at end of file
diff --git a/src/OpenHosta/exec/predict/model/neural_network.py b/src/OpenHosta/exec/predict/model/neural_network.py
index ae61278..65e366f 100644
--- a/src/OpenHosta/exec/predict/model/neural_network.py
+++ b/src/OpenHosta/exec/predict/model/neural_network.py
@@ -4,7 +4,7 @@
 from torch import nn
 
 from .neural_network_types import LayerType, OptimizerAlgorithm, LossFunction, Layer
-from ....utils.torch_nn_utils import pytorch_layer_to_custom, pytorch_loss_to_custom, pytorch_optimizer_to_custom
+from ....utils.torch_nn_utils import pytorch_layer_to_custom, pytorch_loss_to_custom, pytorch_optimizer_to_custom, custom_layer_to_pytorch
 
 
 class NeuralNetwork:
@@ -15,6 +15,7 @@ def __init__(self):
         self.layers: list[Layer] = []
         self.loss_function: Optional[LossFunction] = None
         self.optimizer: Optional[OptimizerAlgorithm] = None
+        self._type = "UNDEFINED"
 
     def add_layer(self, layer: Layer):
         """
@@ -61,6 +62,7 @@ def to_json(self) -> str:
         :rtype: str
         """
         network_dict = {
+            "type": self._type,
             "layers": [
                 layer.to_json()
                 for layer in self.layers
@@ -85,22 +87,33 @@ def from_json(cls, json_str: str) -> 'NeuralNetwork':
         """
         try:
             network_dict = json.loads(json_str)
-            network = cls()
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON: {e}")
 
-            if network_dict.get("loss_function", None) is not None:
-                network.loss_function = LossFunction[network_dict.get("loss_function", None)]
-            else:
-                network.loss_function = None
+        network = cls()
 
-            if network_dict.get("optimizer", None) is not None:
-                network.optimizer = OptimizerAlgorithm[network_dict.get("optimizer", None)]
-            else:
-                network.optimizer = None
+        # Set loss function
+        loss_fn_name = network_dict.get("loss_function")
+        if loss_fn_name:
+            try:
+                network.loss_function = LossFunction[loss_fn_name]
+            except KeyError:
+                raise ValueError(f"Unsupported loss function: {loss_fn_name}")
+
+        # Set optimizer
+        optimizer_name = network_dict.get("optimizer")
+        if optimizer_name:
+            try:
+                network.optimizer = OptimizerAlgorithm[optimizer_name]
+            except KeyError:
+                raise ValueError(f"Unsupported optimizer: {optimizer_name}")
 
-            # Add layers
-            for layer_dict in network_dict.get("layers", []):
+        # Add layers
+        for layer_dict in network_dict.get("layers", []):
+            try:
+                layer_type = LayerType[layer_dict["layer_type"]]
                 layer = Layer(
-                    layer_type=LayerType[layer_dict.get("layer_type", None)],
+                    layer_type=layer_type,
                     in_features=layer_dict.get("in_features"),
                     out_features=layer_dict.get("out_features"),
                     kernel_size=layer_dict.get("kernel_size"),
@@ -109,11 +122,12 @@ def from_json(cls, json_str: str) -> 'NeuralNetwork':
                     dropout=layer_dict.get("dropout")
                 )
                 network.add_layer(layer)
+            except KeyError as e:
+                raise ValueError(f"Missing key in layer definition: {e}")
+            except ValueError as e:
+                raise ValueError(f"Invalid value in layer definition: {e}")
 
-            return network
-
-        except (json.JSONDecodeError, KeyError, ValueError) as e:
-            raise ValueError(f"Invalid JSON configuration: {str(e)}")
+        return network
 
 
     @classmethod
@@ -130,13 +144,13 @@ def from_torch_nn(cls, torch_model: nn.Module, loss_fn=None, optimizer=None) ->
         network = cls()
 
         # Iterating through the PyTorch model's children (layers)
-        for layer in torch_model.children():
-            try:
-                nn_layer = pytorch_layer_to_custom(layer)
+        for layer in torch_model.model.children(): # assuming that the model as a attribut model ! 
+            nn_layer = pytorch_layer_to_custom(layer)
+            if nn_layer is not None:
                 network.add_layer(nn_layer)
-            except ValueError as e:
-                # print(f"Skipping unsupported layer: {e}")
+            else:
                 pass
+                # print(f"Skipping unsupported layer: {e}")
 
         # Set loss function and optimizer if specified
         if loss_fn is not None:
@@ -153,3 +167,19 @@ def from_torch_nn(cls, torch_model: nn.Module, loss_fn=None, optimizer=None) ->
                 print(f"Skipping unsupported optimizer: {e}")
 
         return network
+    
+    def to_pytorch_sequential_model(self) -> nn.Sequential:
+        """
+        Convert the NeuralNetwork instance to a PyTorch nn.Sequential model.
+
+        :return: A PyTorch nn.Sequential model.
+        """
+        layers = []
+        for layer in self.layers:
+            try:
+                pytorch_layer = custom_layer_to_pytorch(layer)
+                if pytorch_layer is not None:
+                    layers.append(pytorch_layer)
+            except ValueError as e:
+                print(f"Skipping unsupported layer: {e}")
+        return nn.Sequential(*layers)
diff --git a/src/OpenHosta/exec/predict/model/neural_network_types.py b/src/OpenHosta/exec/predict/model/neural_network_types.py
index 1baae4d..a6d576b 100644
--- a/src/OpenHosta/exec/predict/model/neural_network_types.py
+++ b/src/OpenHosta/exec/predict/model/neural_network_types.py
@@ -85,8 +85,8 @@ class Layer:
     Initialize a Layer object.
 
     :param layer_type: The type of the layer.
-    :param in_features: Number of _inputs features or channels.
-    :param out_features: Number of _outputs features or channels.
+    :param in_features: Number of inputs features or channels.
+    :param out_features: Number of outputs features or channels.
     :param kernel_size: Size of the kernel/filter.
     :param stride: Stride of the kernel/filter.
     :param padding: Padding added to the _inputs.
diff --git a/src/OpenHosta/exec/predict/predict.py b/src/OpenHosta/exec/predict/predict.py
index c5d7e2f..292b229 100644
--- a/src/OpenHosta/exec/predict/predict.py
+++ b/src/OpenHosta/exec/predict/predict.py
@@ -1,21 +1,23 @@
 import os
-from typing import Union, Optional
+from typing import Union, Optional, Literal
 
 from .dataset.dataset import HostaDataset, SourceType
 from .dataset.oracle import LLMSyntheticDataGenerator
 from .model import HostaModel
 from .model.model_provider import HostaModelProvider
 from .model.neural_network import NeuralNetwork
+from .model.neural_network_types import ArchitectureType
 from .predict_config import PredictConfig
 from .predict_memory import PredictMemory, File
 from ...core.config import Model, DefaultModel
 from ...core.hosta import Hosta, Func
+from ...core.logger import Logger, ANSIColor
 
 
 def predict(
     config: PredictConfig = PredictConfig(),
     oracle: Optional[Union[Model, HostaDataset]] = None,
-    verbose: int = 0 
+    verbose: Union[Literal[0, 1, 2], bool] = 2
 ) -> Union[int, float, bool, str]:
     """
     Predicts a result using an existing model or by creating a new one.
@@ -28,133 +30,168 @@ def predict(
     Returns:
         Model prediction
     """
-    assert config is not None, "Please provide a valid configuration not None"
-    assert verbose is not None and 0 <= verbose <= 2, "Please provide a valid verbose level (0, 1 or 2) default is 0"
 
-    func: Func = getattr(Hosta(), "_infos")
+    assert config is not None, "Please provide a valid configuration not None"
 
+    x = Hosta()
+    func: Func = getattr(x._update_call(), "_infos")
+    
     name = config.name if config and config.name else str(func.f_name)
     base_path = config.path if config and config.path else os.getcwd()
     memory: PredictMemory = PredictMemory.load(base_path=base_path, name=name)
 
-    dataset: Optional[HostaDataset] = None
-    
-    hosta_model: HostaModel = get_hosta_model(memory.architecture, func, config, verbose)
-    if verbose == 2:
-        print(f"[\033[92mArchitecture\033[0m] loaded, type : {type(hosta_model).__name__}")
-        
-    if not load_weights(memory, hosta_model, verbose):
-        train_model(config, memory, hosta_model, dataset, oracle, func, verbose)
+    logger: Logger = Logger(log_file_path=memory.summary.path, verbose=verbose)
+
+    dataset: Optional[HostaDataset] = getattr(func.f_obj, "_dataset", None)
+
+    hosta_model: HostaModel = get_hosta_model(func, memory.architecture, logger, config)
+
+    if not load_weights(x, memory, hosta_model, logger):
+        train_model(config, memory, hosta_model, oracle, func, logger)
     
     if dataset is None:
-        dataset = HostaDataset.from_input(func.f_args, verbose)
+        dataset = HostaDataset.from_input(func.f_args, logger, config.max_tokens, func, memory.dictionary.path)
+        x.attach(func.f_obj, {"_dataset": dataset})
     else:
-        dataset.prepare_inference(func.f_args)
+        dataset.prepare_inference(func.f_args, config.max_tokens, func, memory.dictionary.path)
+
+    if not hasattr(func.f_obj, "_model"):
+        setattr(func, "_model", hosta_model)
+
+    if config.normalize:
+        dataset.normalize_input_inference(memory.normalization)
     torch_prediction = hosta_model.inference(dataset.inference.input)
-    prediction = dataset.decode(torch_prediction, func_f_type=func.f_type[1])
-    if predict is list:
-        return prediction[0]
-    else:
-        return prediction
+    if config.normalize:
+        torch_prediction = dataset.denormalize_output_inference(torch_prediction, memory.normalization)
+    output, prediction = dataset.decode(torch_prediction, func_f_type=func.f_type)
+    logger.log_custom("Prediction", f"{prediction} -> {output}", color=ANSIColor.BLUE_BOLD, level=1)
+    return output
 
 
-def get_hosta_model(architecture_file: File, func: Func, config: Optional[PredictConfig] = None, verbose: int = 0) -> HostaModel:
+def get_hosta_model(func: Func, architecture_file: File, logger: Logger, config: Optional[PredictConfig] = None) -> HostaModel:
     """
     Load or create a new model.
     """
-    architecture: Optional[NeuralNetwork] = None
+    if hasattr(func.f_obj, "_model"):
+        return getattr(func.f_obj, "_model")
+
+    architecture: Optional[NeuralNetwork] = load_architecure(architecture_file, logger)
+
+    model = HostaModelProvider.from_hosta_func(func, config, architecture, architecture_file.path, logger)
+    return model
 
+
+def load_architecure(architecture_file: File, logger: Logger) -> Union[NeuralNetwork, None]:
+    """
+    Load the architecture if it exists.
+    """
     if architecture_file.exist:
-        with open(architecture_file.path, "r") as file:
+        with open(architecture_file.path, 'r', encoding='utf-8') as file:
             json = file.read()
-        architecture = NeuralNetwork.from_json(json)
-        if verbose == 2:
-            print(f"[\033[92mArchitecture\033[0m] found at {architecture_file.path}")
-    else:
-        if verbose == 2:
-            print(f"[\033[93mArchitecture\033[0m] not found, creating one")
-    return HostaModelProvider.from_hosta_func(func, config, architecture, architecture_file.path, verbose)
-
+        logger.log_custom("Architecture", f"found at {architecture_file.path}", color=ANSIColor.BRIGHT_GREEN, level=2)
+        return NeuralNetwork.from_json(json)
+    else :
+        logger.log_custom("Architecture", "not found", color=ANSIColor.BRIGHT_YELLOW, level=2)
+        return None
 
-def load_weights(memory: PredictMemory, hosta_model: HostaModel, verbose :int) -> bool:
+def load_weights(x: Hosta, memory: PredictMemory, hosta_model: HostaModel, logger: Logger) -> bool:
     """
     Load weights if they exist.
     """
-    if memory.weights.exist:
+    # if hasattr(x._infos.f_obj, "_weights_loaded"):
+    #     return True
 
-        if verbose == 2:
-            print(f"[\033[92mWeights\033[0m] found at {memory.weights.path}")
+    if memory.weights.exist:
+        logger.log_custom("Weights", f"found at {memory.weights.path}", color=ANSIColor.BRIGHT_GREEN, level=2)
         hosta_model.init_weights(memory.weights.path)
+        x.attach(x._infos.f_obj, {"_weights_loaded": True})
         return True
 
-    if verbose == 2:
-        print(f"[\033[92mWeights\033[0m] not found generate new ones")
-    return False 
+    logger.log_custom("Weights", "not found", color=ANSIColor.BRIGHT_YELLOW, level=2)
+    return False
 
 
-def train_model(config: PredictConfig, memory: PredictMemory, model: HostaModel, dataset: HostaDataset, oracle: Optional[Union[Model, HostaDataset]], func: Func, verbose: int) -> None:
+def train_model(config: PredictConfig, memory: PredictMemory, model: HostaModel, oracle: Optional[Union[Model, HostaDataset]], func: Func, logger: Logger) -> None:
     """
     Prepare the data and train the model.
     """
     if memory.data.exist:
-        if verbose == 2:
-            print(f"[\033[92mData\033[0m] found at {memory.data.path}")
-        train_set, val_set = HostaDataset.from_data(memory.data.path, batch_size=1, shuffle=True, train_set_size=0.8, verbose=verbose) # verbose will prcess all the example and add it to val_set
+        logger.log_custom("Data", f"found at {memory.data.path}", color=ANSIColor.BRIGHT_GREEN, level=2)
+        dataset = HostaDataset.from_data(memory.data.path, logger=logger)
+        if config.batch_size is None:
+            config.batch_size = max(1, min(16384, int(0.05 * len(dataset.data))))
+
+        train_set, val_set = dataset.to_dataloaders(batch_size=config.batch_size, shuffle=True, train_ratio=0.8)
+
     else:
-        if verbose == 2:
-            print(f"[\033[93mData\033[0m] not processed, preparing data")
-        train_set, val_set = prepare_dataset(config, memory, dataset, func, oracle, verbose)
+        logger.log_custom("Data", "not found", color=ANSIColor.BRIGHT_YELLOW, level=2)
+        train_set, val_set = prepare_dataset(config, memory, func, oracle, model, logger)
+
+    logger.log_custom("Training", f"epochs: {config.epochs}, batch_size: {config.batch_size}, train_set size: {len(train_set)}, val_set size: {len(val_set)}", color=ANSIColor.BRIGHT_YELLOW, level=2)
 
     if config.epochs is None:
         config.epochs = int(2 * len(train_set.dataset) / config.batch_size if config.batch_size != len(train_set.dataset)\
                                 else 2 * len(train_set.dataset))
         assert config.epochs > 0, f"epochs must be greater than 0 now it's {config.epochs}"
-
+    
     model.trainer(train_set, epochs=config.epochs)
 
-    if verbose > 0:
+    if logger.verbose >= 1:
         model.validate(val_set)
-    
+
+
     model.save_weights(memory.weights.path)
 
 
-def prepare_dataset(config: PredictConfig, memory: PredictMemory, dataset: HostaDataset, func: Func, oracle: Optional[Union[Model, HostaDataset]], verbose: int) -> tuple:
+def prepare_dataset(config: PredictConfig, memory: PredictMemory, func: Func, oracle: Optional[Union[Model, HostaDataset]], model: HostaModel, logger: Logger) -> tuple:
     """
     Prepare the dataset for training.
     """
+
+    if config.dataset_path is None:
+        generated_data_path = os.path.join(memory.predict_dir, "generated_data.csv")
+        if os.path.exists(generated_data_path) and os.path.getsize(generated_data_path) > 0:
+            logger.log_custom("Dataset", "no data.json found, but found generated_data.csv, loading it", color=ANSIColor.BRIGHT_GREEN, level=2)
+            config.dataset_path = generated_data_path
+
     if config.dataset_path is not None:
-        if verbose == 2:
-            print(f"[\033[92mDataset\033[0m] found at {config.dataset_path}")
-        dataset = HostaDataset.from_files(config.dataset_path, SourceType.CSV, verbose) # or JSONL jsp comment faire la détection la
+        logger.log_custom("Dataset", f"found at {config.dataset_path}", color=ANSIColor.BRIGHT_GREEN, level=2)
+        dataset = HostaDataset.from_files(path=config.dataset_path, source_type=None, log=logger)
     else :
-        if verbose == 2:
-            print(f"[\033[93mDataset\033[0m] not found, generate data")
-        dataset = generate_data(func, oracle, verbose)
-        dataset.save(os.path.join(memory.predict_dir, "generated_data.csv"), SourceType.CSV)
-        if verbose == 2:
-            print(f"[\033[92mDataset\033[0m] generated!")
-    dataset.encode(max_tokens=10)
-    dataset.tensorify()
-    dataset.save_data(memory.data.path)
+        logger.log_custom("Dataset", "not found, generate data", color=ANSIColor.BRIGHT_YELLOW, level=2)
+        dataset = _generate_data(func, oracle, config, logger)
+        save_path = os.path.join(memory.predict_dir, "generated_data.csv")
+        dataset.save(save_path, SourceType.CSV)
+        logger.log_custom("Dataset", f"generated and saved at {save_path}", color=ANSIColor.BRIGHT_GREEN, level=2)
+
 
     if config.batch_size is None:
-        config.batch_size = int(0.05 * len(dataset.data)) if 0.05 * len(dataset.data) > 1 else len(dataset.data)
-    train_set, val_set = dataset.convert_data(batch_size=config.batch_size, shuffle=True, train_set_size=0.8)
+        config.batch_size = max(1, min(16384, int(0.05 * len(dataset.data))))
 
-    if verbose == 2:
-        print(f"[\033[92mDataset\033[0m] processed and saved at {memory.data.path}")
+    dataset.encode(max_tokens=config.max_tokens, inference=False, func=func, dictionary_path=memory.dictionary.path)
+    if config.normalize:
+        dataset.normalize_data(memory.normalization)
+    dataset.tensorize()
+    dataset.save_data(memory.data.path)
+
+    train_set, val_set = dataset.to_dataloaders(batch_size=config.batch_size, shuffle=True, train_ratio=0.8)
+
+    logger.log_custom("Dataset", f"processed and saved at {memory.data.path}", color=ANSIColor.BRIGHT_GREEN, level=2)
     return train_set, val_set
 
 
-def generate_data(func: Func, oracle: Optional[Union[Model, HostaDataset]], verbose: int) -> HostaDataset:
+def _generate_data(func: Func, oracle: Optional[Union[Model, HostaDataset]], config: PredictConfig, logger: Logger) -> HostaDataset:
     """
     Generate data for training.
     """
+    request_amounts = int(config.generated_data / 100) if config.generated_data > 100 else 1
+
     data = LLMSyntheticDataGenerator.generate_synthetic_data(
         func=func,
-        request_amounts=3,  # TODO: make it a parameter
-        examples_in_req=80,  # TODO: make it a parameter
+        logger=logger,
+        request_amounts=request_amounts,
+        examples_in_req=int(config.generated_data / request_amounts),
         model=oracle if oracle is not None else DefaultModel().get_default_model()
     )
-    return HostaDataset.from_list(data, verbose)
+    return HostaDataset.from_list(data, logger)
 
diff --git a/src/OpenHosta/exec/predict/predict_config.py b/src/OpenHosta/exec/predict/predict_config.py
index 645b7ca..997aea3 100644
--- a/src/OpenHosta/exec/predict/predict_config.py
+++ b/src/OpenHosta/exec/predict/predict_config.py
@@ -1,50 +1,45 @@
 from typing import Optional
 
-from .model.neural_network_types import ArchitectureType
-
 
 class PredictConfig:
     def __init__(self,
-                 model_type: ArchitectureType = None,
                  name: str = None,
                  path: str = None,
-                 version: str = None,
-                 complexity: int = 4,
-                 max_tokens: int = 10,
+                 complexity: int = 5,
+                 growth_rate: float = 1.5,
+                 coef_layers : int = 100,
+                 normalize: bool = False,
                  epochs: Optional[int] = None,
                  batch_size: Optional[int] = None,
-                 learning_rate: Optional[float] = None,
-                 get_loss: Optional[float] = None,
-                 dataset_path: Optional[str] = None
-             ):
-        self.model_type: ArchitectureType = model_type
-
+                 max_tokens: int = 1,
+                 dataset_path: Optional[str] = None,
+                 generated_data: Optional[int] = 100,
+                ):
         self.name: str = name
         self.path: str = path
-        self.version: str = version
-
         self.complexity: int = complexity
-        self.max_tokens: int = max_tokens
-
-        self.batch_size: int = batch_size
+        self.growth_rate: float = growth_rate
+        self.coef_layers: int = coef_layers
+        self.normalize: bool = normalize
         self.epochs: int = epochs
-        self.learning_rate: float = learning_rate
-        self.get_loss: float = get_loss
+        self.batch_size: int = batch_size
+        self.max_tokens: int = max_tokens
         self.dataset_path: str = dataset_path
+        self.generated_data: int = generated_data
 
     def to_json(self):
         return f"""{{
             "name": "{self.name}",
-            "model_type": "{self.model_type}",
-            "weight_path": "{self.path}",
-            "version": "{self.version}",
+            "path": "{self.path}",
             "complexity": {self.complexity},
-            "max_tokens": {self.max_tokens},
+            "growth_rate": {self.growth_rate},
+            "coef_layers": {self.coef_layers},
+            "normalize": {self.normalize},
             "epochs": {self.epochs},
             "batch_size": {self.batch_size},
-            "learning_rate": {self.learning_rate},
-            "get_loss": {self.get_loss},
+            "max_tokens": {self.max_tokens},
             "dataset_path": "{self.dataset_path}",
+            "generated_data": {self.generated_data}
         }}"""
 
     @staticmethod
@@ -53,14 +48,14 @@ def from_json(json_str: str):
         data = json.loads(json_str)
         return PredictConfig(
             name=data["name"],
-            model_type=ArchitectureType(data["model_type"]),
             path=data["path"],
-            version=data["version"],
             complexity=data["complexity"],
-            max_tokens=data["max_tokens"],
+            growth_rate=data["growth_rate"],
+            coef_layers=data["coef_layers"],
+            normalize=data["normalize"],
             epochs=data["epochs"],
             batch_size=data["batch_size"],
-            learning_rate=data["learning_rate"],
-            get_loss=data["get_loss"],
+            max_tokens=data["max_tokens"],
             dataset_path=data["dataset_path"],
+            generated_data=data["generated_data"]
         )
diff --git a/src/OpenHosta/exec/predict/predict_memory.py b/src/OpenHosta/exec/predict/predict_memory.py
index 74ebd04..90469cd 100644
--- a/src/OpenHosta/exec/predict/predict_memory.py
+++ b/src/OpenHosta/exec/predict/predict_memory.py
@@ -1,19 +1,26 @@
 import os
+from dataclasses import dataclass
 from enum import Enum
-from typing import Optional, Dict, NamedTuple
+from typing import Optional, Dict
 
 from ...core.memory import HostaMemory
 
-# 1. Structures de base
-File = NamedTuple("File", [("exist", bool), ("path", str)])
+
+@dataclass
+class File:
+    exist: bool
+    path: str
+
 
 class PredictFileType(Enum):
     """Enumaration for different types of files in the prediction memory."""
     ARCHITECTURE = "model.json"
     WEIGHTS = "weights.pth"
-    DICTIONARY = "dictionary.txt"
+    DICTIONARY = "dictionary.json"
     DATA = "data.json"
     SUMMARY = "summary.txt"
+    NORMALIZATION = "normalization.json"
+
 
 class PredictMemory(HostaMemory):
     """
@@ -22,6 +29,7 @@ class PredictMemory(HostaMemory):
 
     He uses the File structure to store the status of each file and his path.
     """
+
     def __init__(self, base_path: Optional[str] = None, *, name: str = None, **kwargs):
         super().__init__(base_path=base_path, **kwargs)
         if name is None:
@@ -41,11 +49,10 @@ def load(base_path: Optional[str] = None, name: str = None) -> 'PredictMemory':
             PredictMemory instance.
         """
         memory = PredictMemory(base_path=base_path, name=name)
-        memory._initialize_predict_directory
+        memory._initialize_predict_directory()
         memory._check_files()
         return memory
 
-    @property
     def _initialize_predict_directory(self) -> None:
         """
         Initializes the directory and file structure for predictions.
@@ -67,16 +74,25 @@ def _check_files(self) -> None:
             self.files[file_type] = File(exist=exists, path=path)
 
     @property
-    def architecture(self) -> File: return self.files[PredictFileType.ARCHITECTURE]
+    def architecture(self) -> File:
+        return self.files[PredictFileType.ARCHITECTURE]
+
+    @property
+    def weights(self) -> File:
+        return self.files[PredictFileType.WEIGHTS]
 
     @property
-    def weights(self) -> File: return self.files[PredictFileType.WEIGHTS]
+    def data(self) -> File:
+        return self.files[PredictFileType.DATA]
 
     @property
-    def data(self) -> File: return self.files[PredictFileType.DATA]
+    def summary(self) -> File:
+        return self.files[PredictFileType.SUMMARY]
 
     @property
-    def summary(self) -> File: return self.files[PredictFileType.SUMMARY]
+    def dictionary(self) -> File:
+        return self.files[PredictFileType.DICTIONARY]
 
     @property
-    def dictionary(self) -> File: return self.files[PredictFileType.DICTIONARY]
+    def normalization(self) -> File:
+        return self.files[PredictFileType.NORMALIZATION]
diff --git a/src/OpenHosta/utils/torch_nn_utils.py b/src/OpenHosta/utils/torch_nn_utils.py
index 0763130..ea38199 100644
--- a/src/OpenHosta/utils/torch_nn_utils.py
+++ b/src/OpenHosta/utils/torch_nn_utils.py
@@ -1,4 +1,4 @@
-from typing import Union
+from typing import get_origin, Literal, Union
 
 from torch import nn
 from torch import optim
@@ -242,9 +242,10 @@ def custom_optimizer_to_pytorch(optimizer_algorithm: OptimizerAlgorithm, model:
 
 def type_size(data, tokens_size=10):
     """
-    Calculate the _inputs/_outputs size based on the type of the _inputs data.
+    Calculate the inputs/outputs size based on the type of the inputs data.
 
     Parameters:
+        tokens_size: The size of the tokens in the _inputs data.
         data: Can be of type int, float, list, tuple, numpy array, PyTorch tensor, set, dict, or string.
 
     Returns:
@@ -266,7 +267,7 @@ def type_size(data, tokens_size=10):
         return sum(type_size(item) for item in data)
     elif data is dict:
         return sum(type_size(k) + type_size(v) for k, v in data.items())
-    # elif isinstance(data, typing._GenericAlias) and get_origin(data) is Literal:
-    #     return len(data.__args__)
+    elif get_origin(data) is Literal:
+        return len(data.__args__)
     else:
         raise TypeError(f'Unsupported data type: {type(data)}')


Test patch:
diff --git a/tests/unitTests/predict/dataset/empty_dataset.csv b/tests/unitTests/predict/dataset/empty_dataset.csv
new file mode 100644
index 0000000..e69de29
diff --git a/tests/unitTests/predict/dataset/empty_file.csv b/tests/unitTests/predict/dataset/empty_file.csv
new file mode 100644
index 0000000..e69de29
diff --git a/tests/unitTests/test_inspector.py b/tests/unitTests/test_inspector.py
index e0943b6..93e9d0f 100644
--- a/tests/unitTests/test_inspector.py
+++ b/tests/unitTests/test_inspector.py
@@ -64,13 +64,13 @@ class TestAttach:
     def test_OneAttr(self):
         def tmp():
             return
-        HI._attach(tmp, {"nb": 3})
+        HI.attach(tmp, {"nb": 3})
         assert tmp.nb == 3
             
     def test_ManyAttr(self):
         def tmp():
             return
-        HI._attach(tmp, {
+        HI.attach(tmp, {
             "a": 1,
             "b": "hi",
             "c": tmp
@@ -83,15 +83,15 @@ def test_InvalidArgs(self):
         def tmp():
             return
         with pytest.raises(ValueError, match=re.escape("[HostaInspector._attach] Invalid arguments")):
-            HI._attach(tmp, "6")
-            HI._attach("a", 1)
+            HI.attach(tmp, "6")
+            HI.attach("a", 1)
     
     def test_BasicMethod(self):
         class TMP:
             def tmp(self):
                 return
         x = TMP()
-        HI._attach(x.tmp, {"a": 1})
+        HI.attach(x.tmp, {"a": 1})
         assert x.tmp.a == 1
     
     def test_InvalidMethod(self):
@@ -99,10 +99,10 @@ def tmp():
             return
         with patch('inspect.ismethod', return_value=True):
             with pytest.raises(AttributeError, match=re.escape("[HostaInspector._attach] Failed to attach attributs. \"__func__\" attribut is missing.")):
-                HI._attach(tmp, {"a": 1})
+                HI.attach(tmp, {"a": 1})
     
     def test_InvalidCallable(self):
         class TMP:
             pass
         with pytest.raises(AttributeError, match=re.escape(f"[HostaInspector._attach] Failed to attach attributs. Object's type not supported: {type(TMP)}.")):
-            HI._attach(TMP, {"a": 1})
+            HI.attach(TMP, {"a": 1})