llmflowstack 1.1.3__tar.gz → 1.2.0__tar.gz

{llmflowstack-1.1.3 → llmflowstack-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llmflowstack
-Version: 1.1.3
+Version: 1.2.0
 Summary: LLMFlowStack is a framework for training and using LLMs (LLaMA, GPT-OSS, Gemma, ...). Supports DAPT, fine-tuning, and distributed inference. Public fork without institution-specific components.
 Author-email: Gustavo Henrique Ferreira Cruz <gustavohferreiracruz@gmail.com>
 License: MIT
@@ -10,7 +10,6 @@ Requires-Dist: accelerate
 Requires-Dist: bert-score
 Requires-Dist: bitsandbytes
 Requires-Dist: chromadb
-Requires-Dist: colorama
 Requires-Dist: datasets
 Requires-Dist: evaluate
 Requires-Dist: huggingface-hub

{llmflowstack-1.1.3 → llmflowstack-1.2.0}/llmflowstack/__init__.py

@@ -1,9 +1,9 @@
-from .models.Gemma import Gemma3
-from .models.GPT_OSS import GPT_OSS
-from .models.LLaMA3 import LLaMA3
-from .models.LLaMA4 import LLaMA4
-from .models.MedGemma import MedGemma
-from .rag.pipeline import RAGPipeline
+from .decoders.Gemma import Gemma3
+from .decoders.GPT_OSS import GPT_OSS
+from .decoders.LLaMA3 import LLaMA3
+from .decoders.LLaMA4 import LLaMA4
+from .decoders.MedGemma import MedGemma
+from .rag import VectorDatabase
 from .schemas.params import (GenerationBeamsParams, GenerationParams,
                              GenerationSampleParams, TrainParams)
 from .utils.evaluation_methods import text_evaluation
@@ -14,10 +14,10 @@ __all__ = [
   "LLaMA3",
   "LLaMA4",
   "MedGemma",
-  "RAGPipeline",
   "GenerationBeamsParams",
   "GenerationParams",
   "GenerationSampleParams",
   "TrainParams",
-  "text_evaluation"
+  "text_evaluation",
+  "VectorDatabase"
 ]

llmflowstack-1.1.3/llmflowstack/base/base.py → llmflowstack-1.2.0/llmflowstack/decoders/BaseDecoder.py

@@ -1,15 +1,14 @@
 import gc
 import json
-import logging
 import os
 import random
 from abc import ABC, abstractmethod
+from logging import getLogger
 from typing import Any, Literal, cast
 from uuid import uuid4
 
 import numpy as np
 import torch
-from colorama import Fore, Style, init
 from datasets import Dataset
 from torch import Tensor
 from transformers import AutoTokenizer, PreTrainedTokenizerBase
@@ -20,15 +19,15 @@ from trl.trainer.sft_trainer import SFTTrainer
 from llmflowstack.callbacks.log_collector import LogCollectorCallback
 from llmflowstack.schemas.params import GenerationParams, TrainParams
 from llmflowstack.utils.exceptions import MissingEssentialProp
+from llmflowstack.utils.logging import LogLevel
 
 
-class BaseModel(ABC):
+class BaseDecoder(ABC):
 	model = None
 	tokenizer = None
 	_model_id = None
 	model_is_quantized = None
 	seed = None
-	log_level: Literal["INFO", "DEBUG", "WARNING"] = "INFO"
 	stop_token_ids = []
 	question_fields = []
 	answer_fields = []
@@ -37,20 +36,17 @@ class BaseModel(ABC):
 		self,
 		checkpoint: str | None = None,
 		quantization: Literal["4bit", "8bit"] | bool | None = None,
-		seed: int | None = None,
-		log_level: Literal["INFO", "DEBUG", "WARNING"] = "INFO",
+		seed: int | None = None
 	) -> None:
 		if not self.question_fields or not self.answer_fields:
 			raise NotImplementedError("Subclasses must define question_fields and answer_fields.")
 
-		init(autoreset=True)
 		if seed:
 			self._set_seed(seed)
 
 		self._base_model = checkpoint
 
-		self._set_logger(log_level)
-		self.log_level = log_level
+		self.logger = getLogger(f"LLMFlowStack.{self.__class__.__name__}")
 
 		self.tokenizer: PreTrainedTokenizerBase | None = None
 
@@ -61,6 +57,17 @@ class BaseModel(ABC):
 				quantization=quantization
 			)
 	
+	def _log(
+		self,
+		message: str,
+		level: LogLevel = LogLevel.INFO,
+	) -> None:
+		log_func = getattr(self.logger, level.lower(), None)
+		if log_func:
+			log_func(message)
+		else:
+			self.logger.info(message)
+	
 	@abstractmethod
 	def _load_model(
 		self,
@@ -84,7 +91,7 @@ class BaseModel(ABC):
 		quantization: Any
 	) -> None:
 		if self.model:
-			self._log("A model is already loaded. Attempting to reset it.", "WARNING")
+			self._log("A model is already loaded. Attempting to reset it.", LogLevel.WARNING)
 			self.unload_model()
 
 		self._log(f"Loading model on '{checkpoint}'")
@@ -132,42 +139,6 @@ class BaseModel(ABC):
 	) -> None:
 		self._model_id = uuid4()
 
-	def _set_logger(
-		self,
-		level: str
-	) -> None:
-		level_map = {
-			"DEBUG": logging.DEBUG,
-			"INFO": logging.INFO,
-			"WARNING": logging.WARNING,
-			"ERROR": logging.ERROR,
-		}
-		numeric_level = level_map.get(level.upper(), logging.INFO)
-	
-		logging.basicConfig(
-			level=numeric_level,
-			format="%(asctime)s - %(levelname)s - %(message)s"
-		)
-		self.logger = logging.getLogger(__name__)
-
-	def _log(
-		self,
-		info: str,
-		level: Literal["INFO", "WARNING", "ERROR", "DEBUG"] = "INFO"
-	) -> None:
-		if level == "INFO":
-			colored_msg = f"{Fore.GREEN}{info}{Style.RESET_ALL}"
-			self.logger.info(colored_msg)
-		elif level == "WARNING":
-			colored_msg = f"{Fore.YELLOW}{info}{Style.RESET_ALL}"
-			self.logger.warning(colored_msg)
-		elif level == "ERROR":
-			colored_msg = f"{Fore.RED}{info}{Style.RESET_ALL}"
-			self.logger.error(colored_msg)
-		elif level == "DEBUG":
-			colored_msg = f"{Fore.BLUE}{info}{Style.RESET_ALL}"
-			self.logger.debug(colored_msg)
-
 	def _set_seed(
 		self,
 		seed: int
@@ -190,10 +161,10 @@ class BaseModel(ABC):
 		path: str
 	) -> None:
 		if not self.model:
-			self._log("No model to save.", "WARNING")
+			self._log("No model to save.", LogLevel.WARNING)
 			return None
 		if not self.tokenizer:
-			self._log("No tokenizer to save.", "WARNING")
+			self._log("No tokenizer to save.", LogLevel.WARNING)
 			return None
 
 		os.makedirs(path, exist_ok=True)
@@ -299,16 +270,16 @@ class BaseModel(ABC):
 		save_path: str | None = None
 	) -> None:
 		if not self.model:
-			self._log("Could not find a model loaded. Try loading a model first.", "WARNING")
+			self._log("Could not find a model loaded. Try loading a model first.", LogLevel.WARNING)
 			return None
 		if not self.tokenizer:
-			self._log("Could not find a tokenizer loaded. Try loading a tokenizer first.", "WARNING")
+			self._log("Could not find a tokenizer loaded. Try loading a tokenizer first.", LogLevel.WARNING)
 			return None
 
 		self._log("Starting DAPT")
 
 		if self.model_is_quantized:
-			self._log("Cannot DAPT a quantized model.", "WARNING")
+			self._log("Cannot DAPT a quantized model.", LogLevel.WARNING)
 			return None
 		
 		if params is None:
@@ -443,16 +414,16 @@ class BaseModel(ABC):
 		save_path: str | None = None
 	) -> None:
 		if not self.model:
-			self._log("Could not find a model loaded. Try loading a model first.", "WARNING")
+			self._log("Could not find a model loaded. Try loading a model first.", LogLevel.WARNING)
 			return None
 		if not self.tokenizer:
-			self._log("Could not find a tokenizer loaded. Try loading a tokenizer first.", "WARNING")
+			self._log("Could not find a tokenizer loaded. Try loading a tokenizer first.", LogLevel.WARNING)
 			return None
 
 		self._log("Starting fine-tune")
 
 		if self.model_is_quantized:
-			self._log("Cannot fine-tune a quantized model.", "WARNING")
+			self._log("Cannot fine-tune a quantized model.", LogLevel.WARNING)
 			return None
 		
 		if params is None:
@@ -521,8 +492,8 @@ class BaseModel(ABC):
 			self._model_id = None
 			self._log("Reset successfully.")
 		except Exception as e:
-			self._log("Couldn't reset model...", "ERROR")
-			self._log(f"{str(e)}", "DEBUG")
+			self._log("Couldn't reset model...", LogLevel.ERROR)
+			self._log(f"{str(e)}", LogLevel.DEBUG)
 
 	def set_seed(self, seed: int) -> None:
 		self._log(f"Setting seed value {seed}")

{llmflowstack-1.1.3/llmflowstack/models → llmflowstack-1.2.0/llmflowstack/decoders}/GPT_OSS.py

@@ -1,4 +1,3 @@
-import textwrap
 import threading
 from functools import partial
 from time import time
@@ -11,11 +10,12 @@ from transformers import (AutoTokenizer, StoppingCriteriaList,
 from transformers.models.gpt_oss import GptOssForCausalLM
 from transformers.utils.quantization_config import Mxfp4Config
 
-from llmflowstack.base.base import BaseModel
 from llmflowstack.callbacks.stop_on_token import StopOnToken
+from llmflowstack.decoders.BaseDecoder import BaseDecoder
 from llmflowstack.schemas.params import GenerationParams
 from llmflowstack.utils.exceptions import MissingEssentialProp
 from llmflowstack.utils.generation_utils import create_generation_params
+from llmflowstack.utils.logging import LogLevel
 
 
 class GPTOSSInput(TypedDict):
@@ -26,7 +26,7 @@ class GPTOSSInput(TypedDict):
 	reasoning_message: str | None
 	reasoning_level: Literal["Low", "Medium", "High"] | None
 
-class GPT_OSS(BaseModel):
+class GPT_OSS(BaseDecoder):
 	model: GptOssForCausalLM | None = None
 	reasoning_level: Literal["Low", "Medium", "High"] = "Low"
 	question_fields = ["input_text", "developer_message", "system_message"]
@@ -36,14 +36,12 @@ class GPT_OSS(BaseModel):
 		self,
 		checkpoint: str | None = None,
 		quantization: bool | None = None,
-		seed: int | None = None,
-		log_level: Literal["INFO", "DEBUG", "WARNING"] = "INFO",
+		seed: int | None = None
 	) -> None:
 		return super().__init__(
 			checkpoint=checkpoint,
 			quantization=quantization,
-			seed=seed,
-			log_level=log_level
+			seed=seed
 		)
 
 	def _set_generation_stopping_tokens(
@@ -51,7 +49,7 @@ class GPT_OSS(BaseModel):
 		tokens: list[int]
 	) -> None:
 		if not self.tokenizer:
-			self._log("Could not set stop tokens - generation may not work...", "WARNING")
+			self._log("Could not set stop tokens - generation may not work...", LogLevel.WARNING)
 			return None
 		encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
 		particular_tokens = encoding.stop_tokens_for_assistant_actions()
@@ -76,7 +74,7 @@ class GPT_OSS(BaseModel):
 				attn_implementation="eager",
 			)
 		except Exception as _:
-			self._log("Error trying to load the model. Defaulting to load without quantization...", "WARNING")
+			self._log("Error trying to load the model. Defaulting to load without quantization...", LogLevel.WARNING)
 			self.model = GptOssForCausalLM.from_pretrained(
 				checkpoint,
 				dtype="auto",
@@ -119,7 +117,11 @@ class GPT_OSS(BaseModel):
 		if expected_answer:
 			assistant_text += f"<|start|>assistant<|channel|>final<|message|>{expected_answer}<|return|>"
 
-		return textwrap.dedent(f"""{system_text}{developer_text}<|start|>user<|message|>{data["input_text"]}<|end|>{assistant_text}""")
+		return (
+			f"{system_text}{developer_text}"
+			f"<|start|>user<|message|>{data["input_text"]}<|end|>"
+			f"{assistant_text}"
+		)
 
 	def build_input(
 		self,
@@ -154,7 +156,7 @@ class GPT_OSS(BaseModel):
 		params: GenerationParams | None = None
 	) -> str | None:
 		if self.model is None or self.tokenizer is None:
-			self._log("Model or Tokenizer missing", "WARNING")
+			self._log("Model or Tokenizer missing", LogLevel.WARNING)
 			return None
 
 		self._log(f"Processing received input...'")
@@ -222,7 +224,7 @@ class GPT_OSS(BaseModel):
 		params: GenerationParams | None = None
 	) -> Iterator[str]:
 		if self.model is None or self.tokenizer is None:
-			self._log("Model or Tokenizer missing", "WARNING")
+			self._log("Model or Tokenizer missing", LogLevel.WARNING)
 			if False:
 				yield ""
 			return

{llmflowstack-1.1.3/llmflowstack/models → llmflowstack-1.2.0/llmflowstack/decoders}/Gemma.py

@@ -10,12 +10,13 @@ from transformers import (AutoTokenizer, DataCollatorForLanguageModeling,
 from transformers.models.gemma3 import Gemma3ForCausalLM
 from transformers.utils.quantization_config import BitsAndBytesConfig
 
-from llmflowstack.base.base import BaseModel
 from llmflowstack.callbacks.log_collector import LogCollectorCallback
 from llmflowstack.callbacks.stop_on_token import StopOnToken
+from llmflowstack.decoders.BaseDecoder import BaseDecoder
 from llmflowstack.schemas.params import GenerationParams, TrainParams
 from llmflowstack.utils.exceptions import MissingEssentialProp
 from llmflowstack.utils.generation_utils import create_generation_params
+from llmflowstack.utils.logging import LogLevel
 
 
 class Gemma3Input(TypedDict):
@@ -24,7 +25,7 @@ class Gemma3Input(TypedDict):
 	system_message: str | None
 	image_paths: list[str] | None
 
-class Gemma3(BaseModel):
+class Gemma3(BaseDecoder):
 	model: Gemma3ForCausalLM | None = None
 	question_fields = ["input_text", "system_message"]
 	answer_fields = ["expected_answer"]
@@ -33,14 +34,12 @@ class Gemma3(BaseModel):
 		self,
 		checkpoint: str | None = None,
 		quantization: Literal["4bit"] | None = None,
-		seed: int | None = None,
-		log_level: Literal["INFO", "DEBUG", "WARNING"] = "INFO",
+		seed: int | None = None
 	) -> None:
 		return super().__init__(
 			checkpoint=checkpoint,
 			quantization=quantization,
-			seed=seed,
-			log_level=log_level
+			seed=seed
 		)
 
 	def _set_generation_stopping_tokens(
@@ -48,7 +47,7 @@ class Gemma3(BaseModel):
 		tokens: list[int]
 	) -> None:
 		if not self.tokenizer:
-			self._log("Could not set stop tokens - generation may not work...", "WARNING")
+			self._log("Could not set stop tokens - generation may not work...", LogLevel.WARNING)
 			return None
 		particular_tokens = self.tokenizer.encode("<end_of_turn>")
 		self.stop_token_ids = tokens + particular_tokens
@@ -129,16 +128,16 @@ class Gemma3(BaseModel):
 		save_path: str | None = None
 	) -> None:
 		if not self.model:
-			self._log("Could not find a model loaded. Try loading a model first.", "WARNING")
+			self._log("Could not find a model loaded. Try loading a model first.", LogLevel.WARNING)
 			return None
 		if not self.tokenizer:
-			self._log("Could not find a tokenizer loaded. Try loading a tokenizer first.", "WARNING")
+			self._log("Could not find a tokenizer loaded. Try loading a tokenizer first.", LogLevel.WARNING)
 			return None
 
 		self._log("Starting Training")
 
 		if self.model_is_quantized:
-			self._log("Cannot traub a quantized model.", "WARNING")
+			self._log("Cannot traub a quantized model.", LogLevel.WARNING)
 			return None
 		
 		if params is None:
@@ -195,7 +194,7 @@ class Gemma3(BaseModel):
 		save_at_end = True,
 		save_path: str | None = None
 	) -> None:
-		self._log("Only 'dapt' method is available for this class. Redirecting call to it.", "WARNING")
+		self._log("Only 'dapt' method is available for this class. Redirecting call to it.", LogLevel.WARNING)
 		return self.dapt(
 			train_dataset=train_dataset,
 			params=params,
@@ -210,7 +209,7 @@ class Gemma3(BaseModel):
 		params: GenerationParams | None = None,
 	) -> str | None:
 		if self.model is None or self.tokenizer is None:
-			self._log("Model or Tokenizer missing", "WARNING")
+			self._log("Model or Tokenizer missing", LogLevel.WARNING)
 			return None
 
 		self._log(f"Processing received input...'")
@@ -267,7 +266,7 @@ class Gemma3(BaseModel):
 		params: GenerationParams | None = None
 	) -> Iterator[str]:
 		if self.model is None or self.tokenizer is None:
-			self._log("Model or Tokenizer missing", "WARNING")
+			self._log("Model or Tokenizer missing", LogLevel.WARNING)
 			if False:
 				yield ""
 			return

{llmflowstack-1.1.3/llmflowstack/models → llmflowstack-1.2.0/llmflowstack/decoders}/LLaMA3.py

@@ -1,4 +1,3 @@
-import textwrap
 import threading
 from time import time
 from typing import Iterator, Literal, TypedDict, cast
@@ -9,11 +8,12 @@ from transformers import (AutoTokenizer, StoppingCriteriaList,
 from transformers.models.llama import LlamaForCausalLM
 from transformers.utils.quantization_config import BitsAndBytesConfig
 
-from llmflowstack.base.base import BaseModel
 from llmflowstack.callbacks.stop_on_token import StopOnToken
+from llmflowstack.decoders.BaseDecoder import BaseDecoder
 from llmflowstack.schemas.params import GenerationParams
 from llmflowstack.utils.exceptions import MissingEssentialProp
 from llmflowstack.utils.generation_utils import create_generation_params
+from llmflowstack.utils.logging import LogLevel
 
 
 class LLaMA3Input(TypedDict):
@@ -21,7 +21,7 @@ class LLaMA3Input(TypedDict):
 	expected_answer: str | None
 	system_message: str | None
 
-class LLaMA3(BaseModel):
+class LLaMA3(BaseDecoder):
 	model: LlamaForCausalLM | None = None
 	question_fields = ["input_text", "system_message"]
 	answer_fields = ["expected_answer"]
@@ -30,14 +30,12 @@ class LLaMA3(BaseModel):
 		self,
 		checkpoint: str | None = None,
 		quantization: Literal["4bit", "8bit"] | None = None,
-		seed: int | None = None,
-		log_level: Literal["INFO", "DEBUG", "WARNING"] = "INFO",
+		seed: int | None = None
 	) -> None:
 		return super().__init__(
 			checkpoint=checkpoint,
 			quantization=quantization,
-			seed=seed,
-			log_level=log_level
+			seed=seed
 		)
 
 	def _set_generation_stopping_tokens(
@@ -45,7 +43,7 @@ class LLaMA3(BaseModel):
 		tokens: list[int]
 	) -> None:
 		if not self.tokenizer:
-			self._log("Could not set stop tokens - generation may not work...", "WARNING")
+			self._log("Could not set stop tokens - generation may not work...", LogLevel.WARNING)
 			return None
 		particular_tokens = self.tokenizer.encode("<|eot_id|>")
 		self.stop_token_ids = tokens + particular_tokens
@@ -92,7 +90,7 @@ class LLaMA3(BaseModel):
 
 		system_message = data.get("system_message", "")
 
-		return textwrap.dedent(
+		return (
 			f"<|start_header_id|>system<|end_header_id|>{system_message}\n"
 			f"<|eot_id|><|start_header_id|>user<|end_header_id|>{data["input_text"]}\n"
 			f"<|eot_id|><|start_header_id|>assistant<|end_header_id|>{answer}"
@@ -119,7 +117,7 @@ class LLaMA3(BaseModel):
 		params: GenerationParams | None = None
 	) -> str | None:
 		if self.model is None or self.tokenizer is None:
-			self._log("Model or Tokenizer missing", "WARNING")
+			self._log("Model or Tokenizer missing", LogLevel.WARNING)
 			return None
 
 		self.model
@@ -184,7 +182,7 @@ class LLaMA3(BaseModel):
 		params: GenerationParams | None = None
 	) -> Iterator[str]:
 		if self.model is None or self.tokenizer is None:
-			self._log("Model or Tokenizer missing", "WARNING")
+			self._log("Model or Tokenizer missing", LogLevel.WARNING)
 			if False:
 				yield ""
 			return

{llmflowstack-1.1.3/llmflowstack/models → llmflowstack-1.2.0/llmflowstack/decoders}/LLaMA4.py

@@ -1,7 +1,7 @@
 import threading
 from functools import partial
 from time import time
-from typing import Iterator, Literal, TypedDict, cast
+from typing import Iterator, TypedDict, cast
 
 import torch
 from transformers import (AutoTokenizer, DataCollatorForLanguageModeling,
@@ -9,12 +9,13 @@ from transformers import (AutoTokenizer, DataCollatorForLanguageModeling,
                           TrainingArguments)
 from transformers.models.llama4 import Llama4ForCausalLM
 
-from llmflowstack.base.base import BaseModel
 from llmflowstack.callbacks.log_collector import LogCollectorCallback
 from llmflowstack.callbacks.stop_on_token import StopOnToken
+from llmflowstack.decoders.BaseDecoder import BaseDecoder
 from llmflowstack.schemas.params import GenerationParams, TrainParams
 from llmflowstack.utils.exceptions import MissingEssentialProp
 from llmflowstack.utils.generation_utils import create_generation_params
+from llmflowstack.utils.logging import LogLevel
 
 
 class LLaMA4Input(TypedDict):
@@ -22,7 +23,7 @@ class LLaMA4Input(TypedDict):
 	expected_answer: str | None
 	system_message: str | None
 
-class LLaMA4(BaseModel):
+class LLaMA4(BaseDecoder):
 	model: Llama4ForCausalLM | None = None
 	question_fields = ["input_text", "system_message"]
 	answer_fields = ["expected_answer"]
@@ -30,14 +31,12 @@ class LLaMA4(BaseModel):
 	def __init__(
 		self,
 		checkpoint: str | None = None,
-		seed: int | None = None,
-		log_level: Literal["INFO", "DEBUG", "WARNING"] = "INFO",
+		seed: int | None = None
 	) -> None:
 		return super().__init__(
 			checkpoint=checkpoint,
 			quantization=None,
-			seed=seed,
-			log_level=log_level
+			seed=seed
 		)
 
 	def _set_generation_stopping_tokens(
@@ -45,7 +44,7 @@ class LLaMA4(BaseModel):
 		tokens: list[int]
 	) -> None:
 		if not self.tokenizer:
-			self._log("Could not set stop tokens - generation may not work...", "WARNING")
+			self._log("Could not set stop tokens - generation may not work...", LogLevel.WARNING)
 			return None
 		particular_tokens = self.tokenizer.encode("<|eot|>")
 		self.stop_token_ids = tokens + particular_tokens
@@ -84,7 +83,8 @@ class LLaMA4(BaseModel):
 			system_message = f"<|header_start|>system<|header_end|>\n\n{system_message}<|eot|>"
 
 		expected_answer = data.get("expected_answer")
-		answer = f"<|header_start|>assistant<|header_end|>\n\n{expected_answer}<|eot|>" if expected_answer else ""
+		answer = "<|header_start|>assistant<|header_end|>\n\n"
+		answer += f"{expected_answer}<|eot|>" if expected_answer else ""
 
 		return (
 			"<|begin_of_text|>"
@@ -118,16 +118,16 @@ class LLaMA4(BaseModel):
 		save_path: str | None = None
 	) -> None:
 		if not self.model:
-			self._log("Could not find a model loaded. Try loading a model first.", "WARNING")
+			self._log("Could not find a model loaded. Try loading a model first.", LogLevel.WARNING)
 			return None
 		if not self.tokenizer:
-			self._log("Could not find a tokenizer loaded. Try loading a tokenizer first.", "WARNING")
+			self._log("Could not find a tokenizer loaded. Try loading a tokenizer first.", LogLevel.WARNING)
 			return None
 
 		self._log("Starting DAPT")
 
 		if self.model_is_quantized:
-			self._log("Cannot DAPT a quantized model.", "WARNING")
+			self._log("Cannot DAPT a quantized model.", LogLevel.WARNING)
 			return None
 		
 		if params is None:
@@ -184,7 +184,7 @@ class LLaMA4(BaseModel):
 		save_at_end = True,
 		save_path: str | None = None
 	) -> None:
-		self._log("Only 'dapt' method is available for this class. Redirecting call to it.", "WARNING")
+		self._log("Only 'dapt' method is available for this class. Redirecting call to it.", LogLevel.WARNING)
 		return self.dapt(
 			train_dataset=train_dataset,
 			params=params,
@@ -199,7 +199,7 @@ class LLaMA4(BaseModel):
 		params: GenerationParams | None = None
 	) -> str | None:
 		if self.model is None or self.tokenizer is None:
-			self._log("Model or Tokenizer missing", "WARNING")
+			self._log("Model or Tokenizer missing", LogLevel.WARNING)
 			return None
 
 		self.model
@@ -263,7 +263,7 @@ class LLaMA4(BaseModel):
 		params: GenerationParams | None = None
 	) -> Iterator[str]:
 		if self.model is None or self.tokenizer is None:
-			self._log("Model or Tokenizer missing", "WARNING")
+			self._log("Model or Tokenizer missing", LogLevel.WARNING)
 			if False:
 				yield ""
 			return

{llmflowstack-1.1.3/llmflowstack/models → llmflowstack-1.2.0/llmflowstack/decoders}/MedGemma.py

@@ -1,4 +1,3 @@
-import textwrap
 import threading
 from functools import partial
 from time import time
@@ -10,11 +9,12 @@ from transformers import (AutoTokenizer, StoppingCriteriaList,
 from transformers.models.gemma3 import Gemma3ForCausalLM
 from transformers.utils.quantization_config import BitsAndBytesConfig
 
-from llmflowstack.base.base import BaseModel
 from llmflowstack.callbacks.stop_on_token import StopOnToken
+from llmflowstack.decoders.BaseDecoder import BaseDecoder
 from llmflowstack.schemas.params import GenerationParams
 from llmflowstack.utils.exceptions import MissingEssentialProp
 from llmflowstack.utils.generation_utils import create_generation_params
+from llmflowstack.utils.logging import LogLevel
 
 
 class MedGemmaInput(TypedDict):
@@ -22,7 +22,7 @@ class MedGemmaInput(TypedDict):
 	expected_answer: str | None
 	system_message: str | None
 
-class MedGemma(BaseModel):
+class MedGemma(BaseDecoder):
 	model: Gemma3ForCausalLM | None = None
 	can_think = False
 	question_fields = ["input_text", "system_message"]
@@ -32,14 +32,12 @@ class MedGemma(BaseModel):
 		self,
 		checkpoint: str | None = None,
 		quantization: Literal["4bit"] | None = None,
-		seed: int | None = None,
-		log_level: Literal["INFO", "DEBUG", "WARNING"] = "INFO",
+		seed: int | None = None
 	) -> None:
 		return super().__init__(
 			checkpoint=checkpoint,
 			quantization=quantization,
-			seed=seed,
-			log_level=log_level
+			seed=seed
 		)
 
 	def _set_generation_stopping_tokens(
@@ -47,7 +45,7 @@ class MedGemma(BaseModel):
 		tokens: list[int]
 	) -> None:
 		if not self.tokenizer:
-			self._log("Could not set stop tokens - generation may not work...", "WARNING")
+			self._log("Could not set stop tokens - generation may not work...", LogLevel.WARNING)
 			return None
 		particular_tokens = self.tokenizer.encode("<end_of_turn>")
 		self.stop_token_ids = tokens + particular_tokens
@@ -128,7 +126,7 @@ class MedGemma(BaseModel):
 		params: GenerationParams | None = None,
 	) -> str | None:
 		if self.model is None or self.tokenizer is None:
-			self._log("Model or Tokenizer missing", "WARNING")
+			self._log("Model or Tokenizer missing", LogLevel.WARNING)
 			return None
 
 		self._log(f"Processing received input...'")
@@ -196,7 +194,7 @@ class MedGemma(BaseModel):
 		params: GenerationParams | None = None
 	) -> Iterator[str]:
 		if self.model is None or self.tokenizer is None:
-			self._log("Model or Tokenizer missing", "WARNING")
+			self._log("Model or Tokenizer missing", LogLevel.WARNING)
 			if False:
 				yield ""
 			return

llmflowstack-1.2.0/llmflowstack/rag/VectorDatabase.py

@@ -0,0 +1,278 @@
+import gc
+import uuid
+from logging import getLogger
+
+import chromadb
+import chromadb.config
+import torch
+from langchain_chroma import Chroma
+from langchain_core.documents import Document
+from langchain_core.embeddings import Embeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from sentence_transformers import SentenceTransformer
+
+from llmflowstack.utils.exceptions import MissingEssentialProp
+from llmflowstack.utils.logging import LogLevel
+
+
+class EncoderWrapper(Embeddings):
+	def __init__(
+		self,
+		model: SentenceTransformer
+	) -> None:
+		self.model = model
+
+	def embed_documents(
+		self,
+		texts: list[str]
+	) -> list[list[float]]:
+		vectors = self.model.encode(texts, task="retrieval", show_progress_bar=False)
+		return vectors.tolist()
+	
+	def embed_query(
+		self,
+		text: str
+	) -> list[float]:
+		vectors = self.model.encode(text, task="retrieval", show_progress_bar=False)
+		return vectors.tolist()
+
+class VectorDatabase:
+	encoder: SentenceTransformer | None = None
+	collections: dict[str, Chroma] = {}
+
+	def __init__(
+		self,
+		checkpoint: str | None = None,
+		chunk_size: int = 1000,
+		chunk_overlap: int = 200
+	) -> None:
+		self.logger = getLogger(f"LLMFlowStack.{self.__class__.__name__}")
+
+		self.encoder = None
+		if checkpoint:
+			self.load_encoder(
+				checkpoint=checkpoint
+			)
+
+		self.splitter = RecursiveCharacterTextSplitter(
+			chunk_size=chunk_size,
+			chunk_overlap=chunk_overlap,
+			add_start_index=True,
+		)
+	
+	def _log(
+		self,
+		message: str,
+		level: LogLevel = LogLevel.INFO,
+	) -> None:
+		log_func = getattr(self.logger, level.lower(), None)
+		if log_func:
+			log_func(message)
+		else:
+			self.logger.info(message)
+	
+	def load_encoder(
+		self,
+		checkpoint: str
+	) -> None:
+		if self.encoder:
+			self._log("A encoder is already loaded. Attempting to reset it.", LogLevel.WARNING)
+			self.unload_encoder()
+
+		self._log(f"Loading encoder on '{checkpoint}'")
+		self.encoder = SentenceTransformer(
+			checkpoint,
+			trust_remote_code=True
+		)
+
+		self._log("Encoder loaded")
+	
+	def unload_encoder(
+		self
+	) -> None:
+		try:
+			del self.encoder
+			gc.collect()
+			torch.cuda.empty_cache()
+			self.encoder = None
+			self._log("Reset successfully.")
+		except Exception as e:
+			self._log("Couldn't reset encoder...", LogLevel.ERROR)
+			self._log(f"{str(e)}", LogLevel.DEBUG)
+	
+	def get_collection(
+		self,
+		collection_name: str = "rag_memory",
+		persist_directory: str | None = None
+	) -> None:
+		if not self.encoder:
+			raise MissingEssentialProp("Could not find encoder.")
+
+		client_settings = chromadb.config.Settings(
+			anonymized_telemetry=False
+		)
+
+		self.collections[collection_name] = Chroma(
+			collection_name=collection_name,
+			embedding_function=EncoderWrapper(self.encoder),
+			persist_directory=persist_directory,
+			client_settings=client_settings
+		)
+	
+	def validate_collection_name(
+		self,
+		collection_name: str
+	) -> None:
+		if collection_name not in self.collections:
+			raise ValueError("Collection name not found in collection")
+
+	def index_documents(
+		self,
+		collection_name: str,
+		docs: list[Document],
+		ids: list[str],
+		can_split: bool = True,
+	) -> None:
+		self.validate_collection_name(
+			collection_name=collection_name
+		)
+
+		for doc, src_id in zip(docs, ids):
+			if doc.metadata is None:
+				doc.metadata = {}
+			doc.metadata["source_id"] = src_id
+
+		if can_split:
+			splits = self.splitter.split_documents(docs)
+		else:
+			splits = docs
+
+		split_ids = []
+		metadatas = []
+		texts = []
+
+		for i, s in enumerate(splits):
+			src = s.metadata.get("source_id", "unknown")
+			sid = f"{src}_{i}"
+			split_ids.append(sid)
+			metadatas.append(s.metadata.copy())
+			texts.append(s.page_content)
+
+		self.collections[collection_name].add_texts(
+			texts=texts,
+			ids=split_ids,
+			metadatas=metadatas
+		)
+	
+	def create(
+		self,
+		collection_name: str,
+		information: str,
+		other_info: dict[str, str] | None = None,
+		doc_id: str | None = None,
+		should_index: bool = True,
+		can_split: bool = True
+	) -> Document:
+		if other_info is None:
+			other_info = {}
+
+		if doc_id is None:
+			doc_id = str(uuid.uuid4())
+			
+		metadata = {"source_id": doc_id, **other_info}
+		doc = Document(
+			page_content=information,
+			metadata=metadata
+		)
+
+		if should_index:
+			self.index_documents(
+				collection_name=collection_name,
+				docs=[doc],
+				ids=[doc_id],
+				can_split=can_split
+			)
+
+		return doc
+	
+	def update(
+		self,
+		collection_name: str,
+		doc_id: str,
+		new_information: str,
+		other_info: dict[str, str] | None = None
+	) -> Document:
+		self.validate_collection_name(
+			collection_name=collection_name
+		)
+
+		if other_info is None:
+			other_info = {}
+	
+		documents_to_delete = self.collections[collection_name].get(
+			where={
+				"source_id": doc_id
+			}
+		)
+
+		ids_to_delete = documents_to_delete.get("ids", [])
+
+		if ids_to_delete:
+			self.collections[collection_name].delete(ids=ids_to_delete)
+
+		return self.create(
+			collection_name=collection_name,
+			information=new_information,
+			other_info=other_info,
+			doc_id=doc_id
+		)
+	
+	def delete(
+		self,
+		collection_name: str,
+		doc_id: str
+	) -> None:
+		self.validate_collection_name(
+			collection_name=collection_name
+		)
+
+		self.collections[collection_name].delete(ids=[doc_id])
+
+	def rquery(
+		self,
+		collection_name: str,
+		query: str,
+		k: int = 4,
+		filter: dict | None = None
+	) -> list[Document]:
+		self.validate_collection_name(
+			collection_name=collection_name
+		)
+		
+		return self.collections[collection_name].similarity_search(
+			query=query,
+			k=k,
+			filter=filter
+		)
+
+	def query(
+		self,
+		collection_name: str,
+		query: str,
+		k: int = 4,
+		filter: dict | None = None
+	) -> str:
+		self.validate_collection_name(
+			collection_name=collection_name
+		)
+
+		if filter:
+			docs = self.collections[collection_name].similarity_search(
+				query=query,
+				k=k,
+				filter=filter
+			)
+		else:
+			docs = self.collections[collection_name].similarity_search(query, k=k)
+
+		return "\n\n".join(doc.page_content for doc in docs)

llmflowstack-1.2.0/llmflowstack/rag/__init__.py

@@ -0,0 +1,5 @@
+from .VectorDatabase import VectorDatabase
+
+__all__ = [
+  "VectorDatabase"
+]

llmflowstack-1.2.0/llmflowstack/utils/logging.py

@@ -0,0 +1,8 @@
+from enum import Enum
+
+
+class LogLevel(str, Enum):
+  INFO = "info"
+  WARNING = "warning"
+  ERROR = "error"
+  DEBUG = "debug"

{llmflowstack-1.1.3 → llmflowstack-1.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "llmflowstack"
-version = "1.1.3"
+version = "1.2.0"
 authors = [
   { name = "Gustavo Henrique Ferreira Cruz", email = "gustavohferreiracruz@gmail.com" }
 ]
@@ -17,7 +17,6 @@ dependencies = [
     "accelerate",
     "bert-score",
     "bitsandbytes",
-    "colorama",
     "chromadb",
     "datasets",
     "evaluate",

llmflowstack-1.1.3/llmflowstack/rag/__iinit__.py

@@ -1,5 +0,0 @@
-from .pipeline import RAGPipeline
-
-__all__ = [
-  "RAGPipeline"
-]

llmflowstack-1.1.3/llmflowstack/rag/pipeline.py

@@ -1,279 +0,0 @@
-import uuid
-
-import chromadb
-import chromadb.config
-from langchain_chroma import Chroma
-from langchain_core.documents import Document
-from langchain_core.embeddings import Embeddings
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from sentence_transformers import SentenceTransformer
-
-
-class EncoderWrapper(Embeddings):
-	def __init__(
-		self,
-		model: SentenceTransformer
-	) -> None:
-		self.model = model
-
-	def embed_documents(
-		self,
-		texts: list[str]
-	) -> list[list[float]]:
-		vectors = self.model.encode(texts, task="retrieval", show_progress_bar=False)
-		return vectors.tolist()
-	
-	def embed_query(
-		self,
-		text: str
-	) -> list[float]:
-		vectors = self.model.encode(text, task="retrieval", show_progress_bar=False)
-		return vectors.tolist()
-
-class RAGPipeline:
-	"""
-	A modular Retrieval-Augmented Generation (RAG) pipeline for embedding, indexing, and retrieving scientific or textual data using SentenceTransformers and Chroma as a vector store.
-
-	Supports both persistent (disk-based) and transient (in-memory) modes depending on whether `persist_directory` is provided.
-	"""
-	def __init__(
-		self,
-		checkpoint: str,
-		collection_name: str = "rag_memory",
-		persist_directory: str | None = None,
-		chunk_size: int = 1000,
-		chunk_overlap: int = 200
-	) -> None:
-		"""
-		Initializes the RAG pipeline.
-
-		Args:
-			checkpoint (str): Path or name of the SentenceTransformer checkpoint.
-			collection_name (str): Name of the Chroma collection to create or load.
-			persist_directory (str | None): Directory where the vector database is stored. If None, all data is kept in-memory and discarded after the session ends.
-			chunk_size (int): Maximum size (in characters) for text chunks during indexing.
-			chunk_overlap (int): Overlap (in characters) between consecutive text chunks.
-		"""
-		self.encoder = SentenceTransformer(checkpoint, trust_remote_code=True)
-
-		client_settings = chromadb.config.Settings(
-			anonymized_telemetry=False
-		)
-
-		self.collection = Chroma(
-			collection_name=collection_name,
-			embedding_function=EncoderWrapper(self.encoder),
-			persist_directory=persist_directory,
-			client_settings=client_settings
-		)
-
-		self.splitter = RecursiveCharacterTextSplitter(
-			chunk_size=chunk_size,
-			chunk_overlap=chunk_overlap,
-			add_start_index=True,
-		)
-	
-	def index_documents(
-		self,
-		docs: list[Document],
-		ids: list[str],
-		can_split: bool = True
-	) -> None:
-		"""
-		Indexes a list of documents into the Chroma vector store.
-
-		Each document is assigned a unique `source_id` and, optionally, split into smaller chunks for more granular retrieval. Each resulting chunk is embedded and stored with its metadata for later similarity search.
-
-		Args:
-			docs (list[Document]): List of LangChain `Document` objects to index.
-			ids (list[str]): Unique identifiers corresponding to each document.
-			can_split (bool): Whether to split documents into smaller chunks before
-				indexing. Set to False to index each document as a single entry
-				(e.g., for short or self-contained texts).
-
-		Returns:
-			None
-		"""
-		for doc, src_id in zip(docs, ids):
-			if doc.metadata is None:
-				doc.metadata = {}
-			doc.metadata["source_id"] = src_id
-
-		if can_split:
-			splits = self.splitter.split_documents(docs)
-		else:
-			splits = docs
-
-		split_ids = []
-		metadatas = []
-		texts = []
-
-		for i, s in enumerate(splits):
-			src = s.metadata.get("source_id", "unknown")
-			sid = f"{src}_{i}"
-			split_ids.append(sid)
-			metadatas.append(s.metadata.copy())
-			texts.append(s.page_content)
-
-		self.collection.add_texts(
-			texts=texts,
-			ids=split_ids,
-			metadatas=metadatas
-		)
-	
-	def create(
-		self,
-		information: str,
-		other_info: dict[str, str] | None = None,
-		doc_id: str | None = None,
-		should_index: bool = True,
-		can_split: bool = True
-	) -> Document:
-		"""
-		Creates a new `Document` and optionally indexes it in the collection.
-
-		This is a convenience method that wraps both document creation and embedding/indexing in one step. Metadata fields are merged into the document and can include any descriptive information (e.g., title, DOI, year).
-
-		Args:
-			information (str): Main textual content of the document.
-			other_info (dict[str, str] | None): Optional metadata fields to include.
-			doc_id (str | None): Custom document identifier. If None, a UUID is generated.
-			should_index (bool): Whether to immediately add the document to the vector store.
-			can_split (bool): Whether to allow splitting before indexing.
-
-		Returns:
-			Document: The created LangChain `Document` object (indexed if specified).
-		"""
-		if other_info is None:
-			other_info = {}
-
-		if doc_id is None:
-			doc_id = str(uuid.uuid4())
-			
-		metadata = {"source_id": doc_id, **other_info}
-		doc = Document(
-			page_content=information,
-			metadata=metadata
-		)
-
-		if should_index:
-			self.index_documents(
-				docs=[doc],
-				ids=[doc_id],
-				can_split=can_split
-			)
-
-		return doc
-	
-	def update(
-		self,
-		doc_id: str,
-		new_information: str,
-		other_info: dict[str, str] | None = None
-	) -> Document:
-		"""
-		Updates an existing document in the collection with new content and metadata.
-
-		All vector entries associated with the provided `doc_id` are deleted, and a new document is created and re-indexed in their place. This ensures that embeddings remain consistent with the latest text content.
-
-		Args:
-			doc_id (str): Identifier of the document to update.
-			new_information (str): Updated text content for the document.
-			other_info (dict[str, str] | None): Optional new metadata to associate.
-
-		Returns:
-			Document: The newly created (updated) `Document` object.
-		"""
-		if other_info is None:
-			other_info = {}
-	
-		documents_to_delete = self.collection.get(
-			where={
-				"source_id": doc_id
-			}
-		)
-
-		ids_to_delete = documents_to_delete.get("ids", [])
-
-		if ids_to_delete:
-			self.collection.delete(ids=ids_to_delete)
-
-		return self.create(
-			information=new_information,
-			other_info=other_info,
-			doc_id=doc_id
-		)
-	
-	def delete(
-		self,
-		doc_id: str
-	) -> None:
-		"""
-		Deletes all indexed entries associated with a specific document ID.
-
-		Removes all vectors and metadata tied to the provided `doc_id` from the collection. Use this to completely erase a document's content from the indexed database.
-
-		Args:
-			doc_id (str): Identifier of the document to delete.
-
-		Returns:
-			None
-		"""
-		self.collection.delete(ids=[doc_id])
-
-	def rquery(
-		self,
-		query: str,
-		k: int = 4,
-		filter: dict | None = None
-	) -> list[Document]:
-		"""
-		Perform a **raw semantic search** on the collection.
-
-		This method queries the vector store using the provided text query and returns the top-`k` most similar `Document` objects, optionally filtered by metadata.
-
-		Args:
-			query (str): The natural-language query text to embed and search for.
-			k (int, optional): Number of top results to return. Defaults to 4.
-			filter (dict | None, optional): Metadata filter applied during search
-				(e.g., {"type": "article"}). Defaults to None.
-
-		Returns:
-			list[Document]: A list of matching documents sorted by similarity score.
-		"""
-		return self.collection.similarity_search(
-			query=query,
-			k=k,
-			filter=filter
-		)
-
-	def query(
-		self,
-		query: str,
-		k: int = 4,
-		filter: dict | None = None
-	) -> str:
-		"""
-		Perform a **semantic search** and return the combined text content.
-
-		This method wraps `rquery()` and concatenates the retrieved document contents into a single string, suitable for direct use in downstream LLM prompts or text processing.
-
-		Args:
-			query (str): The natural-language query text to search for.
-			k (int, optional): Number of top results to return. Defaults to 4.
-			filter (dict | None, optional): Metadata filter applied during search. If None, all documents are considered.
-
-		Returns:
-			str: A newline-separated string containing the page contents of
-			the retrieved documents.
-		"""
-		if filter:
-			docs = self.collection.similarity_search(
-				query=query,
-				k=k,
-				filter=filter
-			)
-		else:
-			docs = self.collection.similarity_search(query, k=k)
-
-		return "\n\n".join(doc.page_content for doc in docs)