llmflowstack 1.0.2__tar.gz → 1.1.1__tar.gz

{llmflowstack-1.0.2 → llmflowstack-1.1.1}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: llmflowstack
-Version: 1.0.2
-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.
+Version: 1.1.1
+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
 License-File: LICENSE
@@ -14,6 +14,7 @@ Requires-Dist: colorama
 Requires-Dist: datasets
 Requires-Dist: evaluate
 Requires-Dist: huggingface-hub
+Requires-Dist: kernels
 Requires-Dist: langchain-chroma
 Requires-Dist: langchain-community
 Requires-Dist: nltk
@@ -56,18 +57,31 @@ This framework is designed to provide flexibility when working with different op
 
   - [`GPT-OSS 20B`](https://huggingface.co/openai/gpt-oss-20b)
   - [`GPT-OSS 120B`](https://huggingface.co/openai/gpt-oss-120b)
+    > Fine-Tuning, DAPT and Inference Available
 
-- **LLaMA**
+- **LLaMA 3**
 
   - [`LLaMA 3.1 8B - Instruct`](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct)
   - [`LLaMA 3.1 70B - Instruct`](https://huggingface.co/meta-llama/Llama-3.1-70B-Instruct)
   - [`LLaMA 3.3 70B - Instruct`](https://huggingface.co/meta-llama/Llama-3.3-70B-Instruct)
   - [`LLaMA 3.3 405B - Instruct`](https://huggingface.co/meta-llama/Llama-3.1-405B-Instruct)
+    > Fine-Tuning, DAPT and Inference Available
+
+- **LLaMA 4**
+
+  - [`LLaMA 4 Scout - Instruct`](https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct)
+    > DAPT and Inference Available
 
 - **Gemma**
-  - [`MedGemma 27B Text - It`](https://huggingface.co/google/medgemma-27b-text-it)
 
-> Compatibility includes both inference and training (Domain-Adaptive Pre-Training — DAPT — and Supervised Fine-Tuning)
+  - [`Gemma 3 27B - Instruct`](https://huggingface.co/google/gemma-3-27b-it)
+    > DAPT and Inference Available
+
+- **MedGemma**
+  - [`MedGemma 27B Text - Instruct`](https://huggingface.co/google/medgemma-27b-text-it)
+    > Fine-Tuning, DAPT and Inference Available
+
+> Other architectures based on those **may** function correctly.
 
 ---
 

{llmflowstack-1.0.2 → llmflowstack-1.1.1}/README.md

@@ -20,18 +20,31 @@ This framework is designed to provide flexibility when working with different op
 
   - [`GPT-OSS 20B`](https://huggingface.co/openai/gpt-oss-20b)
   - [`GPT-OSS 120B`](https://huggingface.co/openai/gpt-oss-120b)
+    > Fine-Tuning, DAPT and Inference Available
 
-- **LLaMA**
+- **LLaMA 3**
 
   - [`LLaMA 3.1 8B - Instruct`](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct)
   - [`LLaMA 3.1 70B - Instruct`](https://huggingface.co/meta-llama/Llama-3.1-70B-Instruct)
   - [`LLaMA 3.3 70B - Instruct`](https://huggingface.co/meta-llama/Llama-3.3-70B-Instruct)
   - [`LLaMA 3.3 405B - Instruct`](https://huggingface.co/meta-llama/Llama-3.1-405B-Instruct)
+    > Fine-Tuning, DAPT and Inference Available
+
+- **LLaMA 4**
+
+  - [`LLaMA 4 Scout - Instruct`](https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct)
+    > DAPT and Inference Available
 
 - **Gemma**
-  - [`MedGemma 27B Text - It`](https://huggingface.co/google/medgemma-27b-text-it)
 
-> Compatibility includes both inference and training (Domain-Adaptive Pre-Training — DAPT — and Supervised Fine-Tuning)
+  - [`Gemma 3 27B - Instruct`](https://huggingface.co/google/gemma-3-27b-it)
+    > DAPT and Inference Available
+
+- **MedGemma**
+  - [`MedGemma 27B Text - Instruct`](https://huggingface.co/google/medgemma-27b-text-it)
+    > Fine-Tuning, DAPT and Inference Available
+
+> Other architectures based on those **may** function correctly.
 
 ---
 

{llmflowstack-1.0.2 → llmflowstack-1.1.1}/llmflowstack/__init__.py

@@ -1,15 +1,19 @@
-from .models.Gemma import Gemma
+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 .schemas.params import (GenerationBeamsParams, GenerationParams,
                              GenerationSampleParams, TrainParams)
 from .utils.evaluation_methods import text_evaluation
 
 __all__ = [
-  "Gemma",
-  "LLaMA3",
+  "Gemma3",
   "GPT_OSS",
+  "LLaMA3",
+  "LLaMA4",
+  "MedGemma",
   "RAGPipeline",
   "GenerationBeamsParams",
   "GenerationParams",

{llmflowstack-1.0.2 → llmflowstack-1.1.1}/llmflowstack/base/base.py

@@ -1,3 +1,4 @@
+import gc
 import json
 import logging
 import os
@@ -35,7 +36,7 @@ class BaseModel(ABC):
 	def __init__(
 		self,
 		checkpoint: str | None = None,
-		quantization: Literal["8bit", "4bit"] | bool | None = None,
+		quantization: Literal["4bit", "8bit"] | bool | None = None,
 		seed: int | None = None,
 		log_level: Literal["INFO", "DEBUG", "WARNING"] = "INFO",
 	) -> None:
@@ -64,7 +65,8 @@ class BaseModel(ABC):
 	def _load_model(
 		self,
 		checkpoint: str,
-		quantization: Literal["8bit", "4bit"] | bool | None = None
+		*args: Any,
+		**kwargs: Any
 	) -> None:
 		pass
 
@@ -79,7 +81,7 @@ class BaseModel(ABC):
 	def load_checkpoint(
 		self,
 		checkpoint: str,
-		quantization: Literal["8bit", "4bit"] | bool | None = None
+		quantization: Any
 	) -> None:
 		if self.model:
 			self._log("A model is already loaded. Attempting to reset it.", "WARNING")
@@ -223,7 +225,7 @@ class BaseModel(ABC):
 		self,
 		*args: Any,
 		**kwargs: Any
-	) -> str:
+	) -> str | BatchEncoding:
 		pass
 
 	def _tokenize(
@@ -282,7 +284,7 @@ class BaseModel(ABC):
 		output = []
 		for data in dataset:
 			complete_input = self._build_input(
-				**{field: data.get(field) for field in self.question_fields + self.answer_fields}
+				data
 			)
 			output.append(complete_input)
 
@@ -403,13 +405,16 @@ class BaseModel(ABC):
 	def _build_input_for_fine_tune(
 		self,
 		input: dict
-	) -> dict[Literal["partial", "complete"], str]:
+	) -> dict[Literal["partial", "complete"], str | BatchEncoding]:
 		if not self.tokenizer:
 			raise MissingEssentialProp("Could not find tokenizer.")
 
-		partial = self._build_input(**{k: input[k] for k in self.question_fields if k in input})
+		partial = self._build_input({
+			**input,
+			"expected_answer": None
+		})
 
-		complete = self._build_input(**{k: input[k] for k in self.question_fields + self.answer_fields if k in input})
+		complete = self._build_input(input)
 
 		return {
 			"partial": partial,
@@ -508,6 +513,8 @@ class BaseModel(ABC):
 		try:
 			self._log("Trying to reset model...")
 			del self.model
+			gc.collect()
+			torch.cuda.empty_cache()
 			self.model = None
 			self.model_is_quantized = None
 			self.process_id = None

{llmflowstack-1.0.2 → llmflowstack-1.1.1}/llmflowstack/models/GPT_OSS.py

@@ -2,7 +2,7 @@ import textwrap
 import threading
 from functools import partial
 from time import time
-from typing import Any, Generator, Iterator, Literal, TypedDict, cast
+from typing import Iterator, Literal, TypedDict, cast
 
 import torch
 from openai_harmony import HarmonyEncodingName, load_harmony_encoding
@@ -32,6 +32,20 @@ class GPT_OSS(BaseModel):
 	question_fields = ["input_text", "developer_message", "system_message"]
 	answer_fields = ["expected_answer", "reasoning_message"]
 
+	def __init__(
+		self,
+		checkpoint: str | None = None,
+		quantization: bool | None = None,
+		seed: int | None = None,
+		log_level: Literal["INFO", "DEBUG", "WARNING"] = "INFO",
+	) -> None:
+		return super().__init__(
+			checkpoint=checkpoint,
+			quantization=quantization,
+			seed=seed,
+			log_level=log_level
+		)
+
 	def _set_generation_stopping_tokens(
 		self,
 		tokens: list[int]
@@ -46,10 +60,9 @@ class GPT_OSS(BaseModel):
 	def _load_model(
 		self,
 		checkpoint: str,
-		quantization: Literal["8bit", "4bit"] | bool | None = False
+		quantization: bool | None = False
 	) -> None:
 		if quantization:
-			self.model_is_quantized = True
 			quantization_config = Mxfp4Config(dequantize=False)
 		else:
 			quantization_config = Mxfp4Config(dequantize=True)
@@ -70,37 +83,43 @@ class GPT_OSS(BaseModel):
 				device_map="auto",
 				attn_implementation="eager"
 			)
+	
+	def load_checkpoint(
+		self,
+		checkpoint: str,
+		quantization: bool | None = None
+	) -> None:
+		return super().load_checkpoint(checkpoint, quantization)
 
 	def _build_input(
 		self,
-		input_text: str,
-		expected_answer: str | None = None,
-		system_message: str | None = None,
-		reasoning_level: Literal["Low", "Medium", "High"] | None = None,
-		reasoning_message: str | None = None,
-		developer_message: str | None = None
+		data: GPTOSSInput
 	) -> str:
 		if not self.tokenizer:
 			raise MissingEssentialProp("Could not find tokenizer.")
 
-		reasoning = reasoning_level
+		reasoning = data.get("reasoning_level")
 		if reasoning is None:
 			reasoning = self.reasoning_level
 
-		system_text = f"<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.\nKnowledge cutoff: 2024-06\n\nReasoning: {reasoning}\n\n{system_message or ""}# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|>"
+		system_message = data.get("system_message", "")
+		system_text = f"<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.\nKnowledge cutoff: 2024-06\n\nReasoning: {reasoning}\n\n{system_message}# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|>"
 
 		developer_text = ""
+		developer_message = data.get("developer_message", "")
 		if developer_message:
-			developer_text = f"<|start|>developer<|message|># Instructions\n\n{developer_message or ""}<|end|>"
+			developer_text = f"<|start|>developer<|message|># Instructions\n\n{developer_message}<|end|>"
 
 		assistant_text = ""
+		reasoning_message = data.get("reasoning_message", "")
 		if reasoning_message:
 			assistant_text += f"<|start|>assistant<|channel|>analysis<|message|>{reasoning_message}<|end|>"
 
+		expected_answer = data.get("expected_answer", "")
 		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|>{input_text}<|end|>{assistant_text}""")
+		return textwrap.dedent(f"""{system_text}{developer_text}<|start|>user<|message|>{data["input_text"]}<|end|>{assistant_text}""")
 
 	def build_input(
 		self,
@@ -150,15 +169,15 @@ class GPT_OSS(BaseModel):
 
 		model_input = None
 		if isinstance(input, str):
-			model_input = self._build_input(
+			model_input = self.build_input(
 				input_text=input
 			)
+			model_input = self._build_input(
+				data=model_input
+			)
 		else:
 			model_input = self._build_input(
-				input_text=input["input_text"],
-				developer_message=input.get("developer_message", None),
-				system_message=input.get("system_message", None),
-				reasoning_level=input.get("reasoning_level", None)
+				data=input
 			)
 
 		tokenized_input = self._tokenize(model_input)
@@ -217,15 +236,15 @@ class GPT_OSS(BaseModel):
 		self.model.generation_config = generation_params
 
 		if isinstance(input, str):
-			model_input = self._build_input(
+			model_input = self.build_input(
 				input_text=input
 			)
+			model_input = self._build_input(
+				data=model_input
+			)
 		else:
 			model_input = self._build_input(
-				input_text=input["input_text"],
-				developer_message=input.get("developer_message"),
-				system_message=input.get("system_message"),
-				reasoning_level=input.get("reasoning_level")
+				data=input
 			)
 		
 		tokenized_input = self._tokenize(model_input)

llmflowstack-1.1.1/llmflowstack/models/Gemma.py

@@ -0,0 +1,319 @@
+import threading
+from functools import partial
+from time import time
+from typing import Iterator, Literal, TypedDict, cast
+
+import torch
+from transformers import (AutoTokenizer, DataCollatorForLanguageModeling,
+                          StoppingCriteriaList, TextIteratorStreamer, Trainer,
+                          TrainingArguments)
+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.schemas.params import GenerationParams, TrainParams
+from llmflowstack.utils.exceptions import MissingEssentialProp
+from llmflowstack.utils.generation_utils import create_generation_params
+
+
+class Gemma3Input(TypedDict):
+	input_text: str
+	expected_answer: str | None
+	system_message: str | None
+	image_paths: list[str] | None
+
+class Gemma3(BaseModel):
+	model: Gemma3ForCausalLM | None = None
+	question_fields = ["input_text", "system_message"]
+	answer_fields = ["expected_answer"]
+
+	def __init__(
+		self,
+		checkpoint: str | None = None,
+		quantization: Literal["4bit"] | None = None,
+		seed: int | None = None,
+		log_level: Literal["INFO", "DEBUG", "WARNING"] = "INFO",
+	) -> None:
+		return super().__init__(
+			checkpoint=checkpoint,
+			quantization=quantization,
+			seed=seed,
+			log_level=log_level
+		)
+
+	def _set_generation_stopping_tokens(
+		self,
+		tokens: list[int]
+	) -> None:
+		if not self.tokenizer:
+			self._log("Could not set stop tokens - generation may not work...", "WARNING")
+			return None
+		particular_tokens = self.tokenizer.encode("<end_of_turn>")
+		self.stop_token_ids = tokens + particular_tokens
+
+	def _load_model(
+		self,
+		checkpoint: str,
+		quantization: Literal["4bit"] | None = None
+	) -> None:
+		quantization_config = None
+		if quantization == "4bit":
+			quantization_config = BitsAndBytesConfig(
+				load_in_4bit=True
+			)
+
+		self.model = Gemma3ForCausalLM.from_pretrained(
+			checkpoint,
+			quantization_config=quantization_config,
+			dtype="auto",
+			device_map="auto",
+			attn_implementation="eager"
+		)
+	
+	def load_checkpoint(
+		self,
+		checkpoint: str,
+		quantization: Literal['4bit'] | None = None
+	) -> None:
+		return super().load_checkpoint(checkpoint, quantization)
+
+	def _build_input(
+		self,
+		data: Gemma3Input
+	) -> str:
+		if not self.tokenizer:
+			raise MissingEssentialProp("Could not find tokenizer.")
+
+		system_message = data.get("system_message", "")
+		if not system_message:
+			system_message = ""
+
+		if system_message:
+			system_message = f"{system_message}\n"
+
+		expected_answer = data.get("expected_answer")
+		answer = f"{expected_answer}<end_of_turn>" if expected_answer else ""
+	
+		return (
+			f"<start_of_turn>user"
+			f"{system_message}\n{data["input_text"]}<end_of_turn>\n"
+			f"<start_of_turn>model\n"
+			f"{answer}"
+		)
+
+	def build_input(
+		self,
+		input_text: str,
+		system_message: str | None = None,
+		expected_answer: str | None = None,
+		image_paths: list[str] | None = None
+	) -> Gemma3Input:
+		if not self.tokenizer:
+			raise MissingEssentialProp("Could not find tokenizer.")
+
+		return {
+			"input_text": input_text,
+			"system_message": system_message,
+			"expected_answer": expected_answer,
+			"image_paths": image_paths
+		}
+	
+	def dapt(
+		self,
+		train_dataset: list,
+		params: TrainParams | None = None,
+		eval_dataset: list | None = None,
+		save_at_end = True,
+		save_path: str | None = None
+	) -> None:
+		if not self.model:
+			self._log("Could not find a model loaded. Try loading a model first.", "WARNING")
+			return None
+		if not self.tokenizer:
+			self._log("Could not find a tokenizer loaded. Try loading a tokenizer first.", "WARNING")
+			return None
+
+		self._log("Starting Training")
+
+		if self.model_is_quantized:
+			self._log("Cannot traub a quantized model.", "WARNING")
+			return None
+		
+		if params is None:
+			params = TrainParams()
+
+		training_arguments = TrainingArguments(
+			num_train_epochs=params.epochs,
+			learning_rate=params.lr,
+			gradient_accumulation_steps=params.gradient_accumulation,
+			warmup_ratio=params.warmup_ratio,
+			lr_scheduler_type="cosine_with_min_lr",
+			lr_scheduler_kwargs={"min_lr_rate": 0.1},
+			output_dir=None,
+			save_strategy="no",
+			logging_steps=params.logging_steps
+		)
+
+		if self.seed is not None:
+			training_arguments.seed = self.seed
+
+		processed_train_dataset = self._promptfy_dataset_for_dapt(train_dataset)
+		tokenized_train_dataset = self._tokenize_dataset_for_dapt(processed_train_dataset)
+
+		tokenized_eval_dataset = None
+		if eval_dataset:
+			processed_eval_dataset = self._promptfy_dataset_for_dapt(eval_dataset)
+			tokenized_eval_dataset = self._tokenize_dataset_for_dapt(processed_eval_dataset)
+
+		log_callback = LogCollectorCallback()
+
+		trainer = Trainer(
+			model=self.model,
+			train_dataset=tokenized_train_dataset,
+			eval_dataset=tokenized_eval_dataset,
+			args=training_arguments,
+			callbacks=[log_callback],
+			data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False)
+		)
+
+		trainer.train()
+
+		if save_at_end and save_path:
+			self.save_checkpoint(
+				path=save_path
+			)
+
+		self._log("Finished Training")
+	
+	def fine_tune(
+		self,
+		train_dataset: list,
+		params: TrainParams | None = None,
+		eval_dataset: list | None = None,
+		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")
+		return self.dapt(
+			train_dataset=train_dataset,
+			params=params,
+			eval_dataset=eval_dataset,
+			save_at_end=save_at_end,
+			save_path=save_path
+		)
+
+	def generate(
+		self,
+		input: Gemma3Input | str,
+		params: GenerationParams | None = None,
+	) -> str | None:
+		if self.model is None or self.tokenizer is None:
+			self._log("Model or Tokenizer missing", "WARNING")
+			return None
+
+		self._log(f"Processing received input...'")
+
+		if params is None:
+			params = GenerationParams(max_new_tokens=32768)
+		elif params.max_new_tokens is None:
+			params.max_new_tokens = 32768
+
+		generation_params = create_generation_params(params)
+		self.model.generation_config = generation_params
+
+		model_input = None
+		if isinstance(input, str):
+			model_input = self.build_input(
+				input_text=input
+			)
+			model_input = self._build_input(
+				data=model_input
+			)
+		else:
+			model_input = self._build_input(
+				data=input
+			)
+
+		tokenized_input = self._tokenize(model_input)
+		input_ids, attention_mask = tokenized_input
+
+		self.model.eval()
+		self.model.gradient_checkpointing_disable()
+		start = time()
+
+		with torch.no_grad():
+			outputs = self.model.generate(
+				input_ids=input_ids,
+				attention_mask=attention_mask,
+				use_cache=True,
+				eos_token_id=None,
+				stopping_criteria=StoppingCriteriaList([StopOnToken(self.stop_token_ids)])
+			)
+
+		end = time()
+		total_time = end - start
+
+		self._log(f"Response generated in {total_time:.4f} seconds")
+
+		response = outputs[0][input_ids.shape[1]:]
+
+		return self.tokenizer.decode(response, skip_special_tokens=True)
+	
+	def generate_stream(
+		self,
+		input: Gemma3Input | str,
+		params: GenerationParams | None = None
+	) -> Iterator[str]:
+		if self.model is None or self.tokenizer is None:
+			self._log("Model or Tokenizer missing", "WARNING")
+			if False:
+				yield ""
+			return
+
+		if params is None:
+			params = GenerationParams(max_new_tokens=32768)
+		elif params.max_new_tokens is None:
+			params.max_new_tokens = 32768
+
+		generation_params = create_generation_params(params)
+		self.model.generation_config = generation_params
+
+		model_input = None
+		if isinstance(input, str):
+			model_input = self.build_input(
+				input_text=input
+			)
+			model_input = self._build_input(
+				data=model_input
+			)
+		else:
+			model_input = self._build_input(
+				data=input
+			)
+		
+		tokenized_input = self._tokenize(model_input)
+		input_ids, attention_mask = tokenized_input
+
+		streamer = TextIteratorStreamer(
+			cast(AutoTokenizer, self.tokenizer),
+			skip_prompt=True,
+			skip_special_tokens=True
+		)
+
+		generate_fn = partial(
+			self.model.generate,
+			input_ids=input_ids,
+			attention_mask=attention_mask,
+			use_cache=True,
+			eos_token_id=None,
+			streamer=streamer,
+			stopping_criteria=StoppingCriteriaList([StopOnToken(self.stop_token_ids)])
+		)
+
+		thread = threading.Thread(target=generate_fn)
+		thread.start()
+
+		for new_text in streamer:
+			yield new_text

{llmflowstack-1.0.2 → llmflowstack-1.1.1}/llmflowstack/models/LLaMA3.py

@@ -26,6 +26,20 @@ class LLaMA3(BaseModel):
 	question_fields = ["input_text", "system_message"]
 	answer_fields = ["expected_answer"]
 
+	def __init__(
+		self,
+		checkpoint: str | None = None,
+		quantization: Literal["4bit", "8bit"] | None = None,
+		seed: int | None = None,
+		log_level: Literal["INFO", "DEBUG", "WARNING"] = "INFO",
+	) -> None:
+		return super().__init__(
+			checkpoint=checkpoint,
+			quantization=quantization,
+			seed=seed,
+			log_level=log_level
+		)
+
 	def _set_generation_stopping_tokens(
 		self,
 		tokens: list[int]
@@ -39,19 +53,17 @@ class LLaMA3(BaseModel):
 	def _load_model(
 		self,
 		checkpoint: str,
-		quantization: Literal["8bit", "4bit"] | bool | None = None
+		quantization: Literal["4bit", "8bit"] | None = None
 	) -> None:
 		quantization_config = None
 		if quantization == "4bit":
 			quantization_config = BitsAndBytesConfig(
 				load_in_4bit=True
 			)
-			self.model_is_quantized = True
 		if quantization == "8bit":
 			quantization_config = BitsAndBytesConfig(
 				load_in_8bit=True
 			)
-			self.model_is_quantized = True
 
 		self.model = LlamaForCausalLM.from_pretrained(
 			checkpoint,
@@ -60,21 +72,29 @@ class LLaMA3(BaseModel):
 			device_map="auto",
 			attn_implementation="eager"
 		)
+	
+	def load_checkpoint(
+		self,
+		checkpoint: str,
+		quantization: Literal['4bit', "8bit"] | None = None
+	) -> None:
+		return super().load_checkpoint(checkpoint, quantization)
 
 	def _build_input(
 		self,
-		input_text: str,
-		expected_answer: str | None = None,
-		system_message: str | None = None
+		data: LLaMA3Input
 	) -> str:
 		if not self.tokenizer:
 			raise MissingEssentialProp("Could not find tokenizer.")
 
+		expected_answer = data.get("expected_answer")
 		answer = f"{expected_answer}{self.tokenizer.eos_token}" if expected_answer else ""
 
+		system_message = data.get("system_message", "")
+
 		return textwrap.dedent(
-			f"<|start_header_id|>system<|end_header_id|>{system_message or ""}\n"
-			f"<|eot_id|><|start_header_id|>user<|end_header_id|>{input_text}\n"
+			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}"
 		)
 
@@ -120,13 +140,15 @@ class LLaMA3(BaseModel):
 
 		model_input = None
 		if isinstance(input, str):
-			model_input = self._build_input(
+			model_input = self.build_input(
 				input_text=input
 			)
+			model_input = self._build_input(
+				data=model_input
+			)
 		else:
 			model_input = self._build_input(
-				input_text=input["input_text"],
-				system_message=input.get("system_message", "")
+				data=input
 			)
 
 		tokenized_input = self._tokenize(model_input)
@@ -175,14 +197,17 @@ class LLaMA3(BaseModel):
 		generation_params = create_generation_params(params)
 		self.model.generation_config = generation_params
 
+		model_input = None
 		if isinstance(input, str):
-			model_input = self._build_input(
+			model_input = self.build_input(
 				input_text=input
 			)
+			model_input = self._build_input(
+				data=model_input
+			)
 		else:
 			model_input = self._build_input(
-				input_text=input["input_text"],
-				system_message=input.get("system_message")
+				data=input
 			)
 		
 		tokenized_input = self._tokenize(model_input)

llmflowstack-1.1.1/llmflowstack/models/LLaMA4.py

@@ -0,0 +1,317 @@
+import threading
+from functools import partial
+from time import time
+from typing import Iterator, Literal, TypedDict, cast
+
+import torch
+from transformers import (AutoTokenizer, DataCollatorForLanguageModeling,
+                          StoppingCriteriaList, TextIteratorStreamer, Trainer,
+                          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.schemas.params import GenerationParams, TrainParams
+from llmflowstack.utils.exceptions import MissingEssentialProp
+from llmflowstack.utils.generation_utils import create_generation_params
+
+
+class LLaMA4Input(TypedDict):
+	input_text: str
+	expected_answer: str | None
+	system_message: str | None
+	image_paths: list[str] | None
+
+class LLaMA4(BaseModel):
+	model: Llama4ForCausalLM | None = None
+	question_fields = ["input_text", "system_message"]
+	answer_fields = ["expected_answer"]
+
+	def __init__(
+		self,
+		checkpoint: str | None = None,
+		seed: int | None = None,
+		log_level: Literal["INFO", "DEBUG", "WARNING"] = "INFO",
+	) -> None:
+		return super().__init__(
+			checkpoint=checkpoint,
+			quantization=None,
+			seed=seed,
+			log_level=log_level
+		)
+
+	def _set_generation_stopping_tokens(
+		self,
+		tokens: list[int]
+	) -> None:
+		if not self.tokenizer:
+			self._log("Could not set stop tokens - generation may not work...", "WARNING")
+			return None
+		particular_tokens = self.tokenizer.encode("<|eot|>")
+		self.stop_token_ids = tokens + particular_tokens
+
+	def _load_model(
+		self,
+		checkpoint: str,
+		quantization: None = None
+	) -> None:
+		self.model = Llama4ForCausalLM.from_pretrained(
+			checkpoint,
+			dtype="auto",
+			device_map="auto",
+			attn_implementation="eager"
+		)
+	
+	def load_checkpoint(
+		self,
+		checkpoint: str,
+		quantization: None = None
+	) -> None:
+		return super().load_checkpoint(checkpoint, quantization)
+
+	def _build_input(
+		self,
+		data: LLaMA4Input
+	) -> str:
+		if not self.tokenizer:
+			raise MissingEssentialProp("Could not find tokenizer.")
+
+		system_message = data.get("system_message", "")
+		if not system_message:
+			system_message = ""
+
+		if system_message:
+			system_message = f"{system_message}\n"
+
+		expected_answer = data.get("expected_answer")
+		answer = f"{expected_answer}<end_of_turn>" if expected_answer else ""
+	
+		return (
+			f"<start_of_turn>user"
+			f"{system_message}\n{data["input_text"]}<end_of_turn>\n"
+			f"<start_of_turn>model\n"
+			f"{answer}"
+		)
+
+	def build_input(
+		self,
+		input_text: str,
+		system_message: str | None = None,
+		expected_answer: str | None = None,
+		image_paths: list[str] | None = None
+	) -> LLaMA4Input:
+		if not self.tokenizer:
+			raise MissingEssentialProp("Could not find tokenizer.")
+
+		return {
+			"input_text": input_text,
+			"system_message": system_message,
+			"expected_answer": expected_answer,
+			"image_paths": image_paths
+		}
+	
+	def dapt(
+		self,
+		train_dataset: list,
+		params: TrainParams | None = None,
+		eval_dataset: list | None = None,
+		save_at_end = True,
+		save_path: str | None = None
+	) -> None:
+		if not self.model:
+			self._log("Could not find a model loaded. Try loading a model first.", "WARNING")
+			return None
+		if not self.tokenizer:
+			self._log("Could not find a tokenizer loaded. Try loading a tokenizer first.", "WARNING")
+			return None
+
+		self._log("Starting DAPT")
+
+		if self.model_is_quantized:
+			self._log("Cannot DAPT a quantized model.", "WARNING")
+			return None
+		
+		if params is None:
+			params = TrainParams()
+
+		training_arguments = TrainingArguments(
+			num_train_epochs=params.epochs,
+			learning_rate=params.lr,
+			gradient_accumulation_steps=params.gradient_accumulation,
+			warmup_ratio=params.warmup_ratio,
+			lr_scheduler_type="cosine_with_min_lr",
+			lr_scheduler_kwargs={"min_lr_rate": 0.1},
+			output_dir=None,
+			save_strategy="no",
+			logging_steps=params.logging_steps
+		)
+
+		if self.seed is not None:
+			training_arguments.seed = self.seed
+
+		processed_train_dataset = self._promptfy_dataset_for_dapt(train_dataset)
+		tokenized_train_dataset = self._tokenize_dataset_for_dapt(processed_train_dataset)
+
+		tokenized_eval_dataset = None
+		if eval_dataset:
+			processed_eval_dataset = self._promptfy_dataset_for_dapt(eval_dataset)
+			tokenized_eval_dataset = self._tokenize_dataset_for_dapt(processed_eval_dataset)
+
+		log_callback = LogCollectorCallback()
+
+		trainer = Trainer(
+			model=self.model,
+			train_dataset=tokenized_train_dataset,
+			eval_dataset=tokenized_eval_dataset,
+			args=training_arguments,
+			callbacks=[log_callback],
+			data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False)
+		)
+
+		trainer.train()
+
+		if save_at_end and save_path:
+			self.save_checkpoint(
+				path=save_path
+			)
+
+		self._log("Finished DAPT")
+	
+	def fine_tune(
+		self,
+		train_dataset: list,
+		params: TrainParams | None = None,
+		eval_dataset: list | None = None,
+		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")
+		return self.dapt(
+			train_dataset=train_dataset,
+			params=params,
+			eval_dataset=eval_dataset,
+			save_at_end=save_at_end,
+			save_path=save_path
+		)
+
+	def generate(
+		self,
+		input: LLaMA4Input | str,
+		params: GenerationParams | None = None
+	) -> str | None:
+		if self.model is None or self.tokenizer is None:
+			self._log("Model or Tokenizer missing", "WARNING")
+			return None
+
+		self.model
+
+		self._log(f"Processing received input...'")
+
+		if params is None:
+			params = GenerationParams(max_new_tokens=32768)
+		elif params.max_new_tokens is None:
+			params.max_new_tokens = 32768
+
+		generation_params = create_generation_params(params)
+		self.model.generation_config = generation_params
+
+		if params:
+			generation_params = create_generation_params(params)
+			self.model.generation_config = generation_params
+
+		model_input = None
+		if isinstance(input, str):
+			model_input = self.build_input(
+				input_text=input
+			)
+			model_input = self._build_input(
+				data=model_input
+			)
+		else:
+			model_input = self._build_input(
+				data=input
+			)
+
+		tokenized_input = self._tokenize(model_input)
+		input_ids, attention_mask = tokenized_input
+
+		self.model.eval()
+		self.model.gradient_checkpointing_disable()
+
+		start = time()
+
+		with torch.no_grad():
+			outputs = self.model.generate(
+				input_ids=input_ids,
+				attention_mask=attention_mask,
+				use_cache=True,
+				eos_token_id=None,
+				stopping_criteria=StoppingCriteriaList([StopOnToken(self.stop_token_ids)])
+			)
+
+		end = time()
+		total_time = end - start
+
+		self._log(f"Response generated in {total_time:.4f} seconds")
+
+		response = outputs[0][input_ids.shape[1]:]
+
+		return self.tokenizer.decode(response, skip_special_tokens=True)
+	
+	def generate_stream(
+		self,
+		input: LLaMA4Input | str,
+		params: GenerationParams | None = None
+	) -> Iterator[str]:
+		if self.model is None or self.tokenizer is None:
+			self._log("Model or Tokenizer missing", "WARNING")
+			if False:
+				yield ""
+			return
+		
+		if params is None:
+			params = GenerationParams(max_new_tokens=32768)
+		elif params.max_new_tokens is None:
+			params.max_new_tokens = 32768
+
+		generation_params = create_generation_params(params)
+		self.model.generation_config = generation_params
+
+		model_input = None
+		if isinstance(input, str):
+			model_input = self.build_input(
+				input_text=input
+			)
+			model_input = self._build_input(
+				data=model_input
+			)
+		else:
+			model_input = self._build_input(
+				data=input
+			)
+		
+		tokenized_input = self._tokenize(model_input)
+		input_ids, attention_mask = tokenized_input
+
+		streamer = TextIteratorStreamer(
+			cast(AutoTokenizer, self.tokenizer),
+			skip_prompt=True,
+			skip_special_tokens=True
+		)
+
+		generate_fn = partial(
+			self.model.generate,
+			input_ids=input_ids,
+			attention_mask=attention_mask,
+			use_cache=True,
+			eos_token_id=None,
+			streamer=streamer,
+			stopping_criteria=StoppingCriteriaList([StopOnToken(self.stop_token_ids)])
+		)
+
+		thread = threading.Thread(target=generate_fn)
+		thread.start()
+
+		for new_text in streamer:
+			yield new_text

llmflowstack-1.0.2/llmflowstack/models/Gemma.py → llmflowstack-1.1.1/llmflowstack/models/MedGemma.py

@@ -17,17 +17,31 @@ from llmflowstack.utils.exceptions import MissingEssentialProp
 from llmflowstack.utils.generation_utils import create_generation_params
 
 
-class GemmaInput(TypedDict):
+class MedGemmaInput(TypedDict):
 	input_text: str
 	expected_answer: str | None
 	system_message: str | None
 
-class Gemma(BaseModel):
+class MedGemma(BaseModel):
 	model: Gemma3ForCausalLM | None = None
 	can_think = False
 	question_fields = ["input_text", "system_message"]
 	answer_fields = ["expected_answer"]
 
+	def __init__(
+		self,
+		checkpoint: str | None = None,
+		quantization: Literal["4bit"] | None = None,
+		seed: int | None = None,
+		log_level: Literal["INFO", "DEBUG", "WARNING"] = "INFO",
+	) -> None:
+		return super().__init__(
+			checkpoint=checkpoint,
+			quantization=quantization,
+			seed=seed,
+			log_level=log_level
+		)
+
 	def _set_generation_stopping_tokens(
 		self,
 		tokens: list[int]
@@ -41,17 +55,13 @@ class Gemma(BaseModel):
 	def _load_model(
 		self,
 		checkpoint: str,
-		quantization: Literal["8bit", "4bit"] | bool | None = None
+		quantization: Literal["4bit"] | None = None
 	) -> None:
 		quantization_config = None
 		if quantization == "4bit":
 			quantization_config = BitsAndBytesConfig(
 				load_in_4bit=True
 			)
-		if quantization == "8bit":
-			quantization_config = BitsAndBytesConfig(
-				load_in_8bit=True
-			)
 
 		self.model = Gemma3ForCausalLM.from_pretrained(
 			checkpoint,
@@ -60,16 +70,22 @@ class Gemma(BaseModel):
 			device_map="auto",
 			attn_implementation="eager"
 		)
+	
+	def load_checkpoint(
+		self,
+		checkpoint: str,
+		quantization:  Literal["4bit"] | None = None
+	) -> None:
+		return super().load_checkpoint(checkpoint, quantization)
 
 	def _build_input(
 		self,
-		input_text: str,
-		expected_answer: str | None = None,
-		system_message: str | None = None
+		data: MedGemmaInput
 	) -> str:
 		if not self.tokenizer:
 			raise MissingEssentialProp("Could not find tokenizer.")
 
+		system_message = data.get("system_message", "")
 		if not system_message:
 			system_message = ""
 		if self.can_think:
@@ -78,11 +94,12 @@ class Gemma(BaseModel):
 		if system_message:
 			system_message = f"{system_message}\n"
 
+		expected_answer = data.get("expected_answer")
 		answer = f"{expected_answer}<end_of_turn>" if expected_answer else ""
 	
-		return textwrap.dedent(
+		return (
 			f"<start_of_turn>user"
-			f"{system_message}\n{input_text}<end_of_turn>\n"
+			f"{system_message}\n{data["input_text"]}<end_of_turn>\n"
 			f"<start_of_turn>model\n"
 			f"{answer}"
 		)
@@ -92,7 +109,7 @@ class Gemma(BaseModel):
 		input_text: str,
 		expected_answer: str | None = None,
 		system_message: str | None = None
-	) -> GemmaInput:
+	) -> MedGemmaInput:
 		if not self.tokenizer:
 			raise MissingEssentialProp("Could not find tokenizer.")
 
@@ -107,7 +124,7 @@ class Gemma(BaseModel):
 
 	def generate(
 		self,
-		input: GemmaInput | str,
+		input: MedGemmaInput | str,
 		params: GenerationParams | None = None,
 	) -> str | None:
 		if self.model is None or self.tokenizer is None:
@@ -126,17 +143,18 @@ class Gemma(BaseModel):
 
 		model_input = None
 		if isinstance(input, str):
-			model_input = self._build_input(
+			model_input = self.build_input(
 				input_text=input
 			)
+			model_input = self._build_input(
+				data=model_input
+			)
 		else:
 			model_input = self._build_input(
-				input_text=input["input_text"],
-				system_message=input["system_message"]
+				data=input
 			)
 
 		tokenized_input = self._tokenize(model_input)
-
 		input_ids, attention_mask = tokenized_input
 
 		self.model.eval()
@@ -174,7 +192,7 @@ class Gemma(BaseModel):
 	
 	def generate_stream(
 		self,
-		input: GemmaInput | str,
+		input: MedGemmaInput | str,
 		params: GenerationParams | None = None
 	) -> Iterator[str]:
 		if self.model is None or self.tokenizer is None:
@@ -191,14 +209,17 @@ class Gemma(BaseModel):
 		generation_params = create_generation_params(params)
 		self.model.generation_config = generation_params
 
+		model_input = None
 		if isinstance(input, str):
-			model_input = self._build_input(
+			model_input = self.build_input(
 				input_text=input
 			)
+			model_input = self._build_input(
+				data=model_input
+			)
 		else:
 			model_input = self._build_input(
-				input_text=input["input_text"],
-				system_message=input.get("system_message")
+				data=input
 			)
 		
 		tokenized_input = self._tokenize(model_input)

llmflowstack-1.1.1/llmflowstack/models/__init__.py

@@ -0,0 +1,13 @@
+from .Gemma import Gemma3
+from .GPT_OSS import GPT_OSS
+from .LLaMA3 import LLaMA3
+from .LLaMA4 import LLaMA4
+from .MedGemma import MedGemma
+
+__all__ = [
+  "Gemma3",
+  "GPT_OSS",
+  "LLaMA3",
+  "LLaMA4",
+  "MedGemma"
+]

{llmflowstack-1.0.2 → llmflowstack-1.1.1}/llmflowstack/rag/pipeline.py

@@ -1,5 +1,7 @@
 import uuid
 
+import chromadb
+import chromadb.config
 from langchain_chroma import Chroma
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
@@ -38,10 +40,15 @@ class RAGPipeline:
 		
 		self.encoder = SentenceTransformer(checkpoint, trust_remote_code=True)
 
+		client_settings = chromadb.config.Settings(
+			anonymized_telemetry=False
+		)
+
 		self.vector_store = Chroma(
 			collection_name=collection_name,
 			embedding_function=EncoderWrapper(self.encoder),
-			persist_directory=persist_directory
+			persist_directory=persist_directory,
+			client_settings=client_settings
 		)
 
 		self.splitter = RecursiveCharacterTextSplitter(

{llmflowstack-1.0.2 → llmflowstack-1.1.1}/pyproject.toml

@@ -4,11 +4,11 @@ build-backend = "hatchling.build"
 
 [project]
 name = "llmflowstack"
-version = "1.0.2"
+version = "1.1.1"
 authors = [
   { name = "Gustavo Henrique Ferreira Cruz", email = "gustavohferreiracruz@gmail.com" }
 ]
-description = "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."
+description = "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."
 readme = "README.md"
 requires-python = ">=3.12"
 license = {text = "MIT"}
@@ -22,6 +22,7 @@ dependencies = [
     "datasets",
     "evaluate",
     "huggingface-hub",
+    "kernels",
     "langchain-chroma",
     "langchain_community",
     "nltk",

llmflowstack-1.0.2/llmflowstack/models/__init__.py

@@ -1,9 +0,0 @@
-from .Gemma import Gemma
-from .GPT_OSS import GPT_OSS
-from .LLaMA3 import LLaMA3
-
-__all__ = [
-  "Gemma",
-  "GPT_OSS",
-  "LLaMA3"
-]