llmflowstack 1.1.4__tar.gz → 1.2.1__tar.gz

{llmflowstack-1.1.4 → llmflowstack-1.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llmflowstack
-Version: 1.1.4
+Version: 1.2.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
@@ -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.4 → llmflowstack-1.2.1}/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.4/llmflowstack/base/base.py → llmflowstack-1.2.1/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.4/llmflowstack/models → llmflowstack-1.2.1/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):
@@ -24,11 +24,11 @@ class GPTOSSInput(TypedDict):
 	developer_message: str | None
 	expected_answer: str | None
 	reasoning_message: str | None
-	reasoning_level: Literal["Low", "Medium", "High"] | None
+	reasoning_level: Literal["Low", "Medium", "High", "Off"] | None
 
-class GPT_OSS(BaseModel):
+class GPT_OSS(BaseDecoder):
 	model: GptOssForCausalLM | None = None
-	reasoning_level: Literal["Low", "Medium", "High"] = "Low"
+	reasoning_level: Literal["Low", "Medium", "High", "Off"] = "Low"
 	question_fields = ["input_text", "developer_message", "system_message"]
 	answer_fields = ["expected_answer", "reasoning_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",
@@ -104,6 +102,8 @@ class GPT_OSS(BaseModel):
 
 		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|>"
+		if reasoning == "Off":
+			system_text = f"<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.\nKnowledge cutoff: 2024-06\n\n{system_message}# Valid channels: final. Channel must be included for every message.<|end|>"
 
 		developer_text = ""
 		developer_message = data.get("developer_message", "")
@@ -119,7 +119,14 @@ 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}""")
+		if not expected_answer and reasoning == "Off":
+			assistant_text += "<|start|>assistant<|channel|>final<|message|>"
+
+		return (
+			f"{system_text}{developer_text}"
+			f"<|start|>user<|message|>{data["input_text"]}<|end|>"
+			f"{assistant_text}"
+		)
 
 	def build_input(
 		self,
@@ -128,7 +135,7 @@ class GPT_OSS(BaseModel):
 		developer_message: str | None = None,
 		expected_answer: str | None = None,
 		reasoning_message: str | None = None,
-		reasoning_level: Literal["Low", "Medium", "High"] | None = None
+		reasoning_level: Literal["Low", "Medium", "High", "Off"] | None = None
 	) -> GPTOSSInput:
 		if not self.tokenizer:
 			raise MissingEssentialProp("Could not find tokenizer.")
@@ -144,7 +151,7 @@ class GPT_OSS(BaseModel):
 
 	def set_reasoning_level(
 		self,
-		level: Literal["Low", "Medium", "High"]
+		level: Literal["Low", "Medium", "High", "Off"]
 	) -> None:
 		self.reasoning_level = level
 
@@ -154,7 +161,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,11 +229,13 @@ 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
 		
+		self._log(f"Processing received input...'")
+		
 		if params is None:
 			params = GenerationParams(max_new_tokens=32768)
 		elif params.max_new_tokens is None:
@@ -266,19 +275,26 @@ class GPT_OSS(BaseModel):
 			stopping_criteria=StoppingCriteriaList([StopOnToken(self.stop_token_ids)])
 		)
 
+		start = time()
+
 		thread = threading.Thread(target=generate_fn)
 		thread.start()
 
-		done_thinking = False
+		done_thinking = self.reasoning_level == "Off"
 		buffer = ""
 
 		for new_text in streamer:
 			buffer += new_text
 
-			if "final" in buffer:
+			if "final" in buffer and not done_thinking:
 				done_thinking = True
 				buffer = buffer.split("final", 1)[1]
 			
 			if done_thinking:
 				yield buffer
-				buffer = ""
+				buffer = ""
+		
+		end = time()
+		total_time = end - start
+
+		self._log(f"Response generated in {total_time:.4f} seconds")

{llmflowstack-1.1.4/llmflowstack/models → llmflowstack-1.2.1/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,10 +266,12 @@ 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
+		
+		self._log(f"Processing received input...'")
 
 		if params is None:
 			params = GenerationParams(max_new_tokens=32768)
@@ -312,8 +313,15 @@ class Gemma3(BaseModel):
 			stopping_criteria=StoppingCriteriaList([StopOnToken(self.stop_token_ids)])
 		)
 
+		start = time()
+
 		thread = threading.Thread(target=generate_fn)
 		thread.start()
 
 		for new_text in streamer:
-			yield new_text
+			yield new_text
+		
+		end = time()
+		total_time = end - start
+
+		self._log(f"Response generated in {total_time:.4f} seconds")

{llmflowstack-1.1.4/llmflowstack/models → llmflowstack-1.2.1/llmflowstack/decoders}/LLaMA3.py

@@ -1,5 +1,5 @@
-import textwrap
 import threading
+from functools import partial
 from time import time
 from typing import Iterator, Literal, TypedDict, cast
 
@@ -9,11 +9,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 +22,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 +31,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 +44,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 +91,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 +118,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,11 +183,13 @@ 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
 		
+		self._log(f"Processing received input...'")
+		
 		if params is None:
 			params = GenerationParams(max_new_tokens=8192)
 		elif params.max_new_tokens is None:
@@ -219,20 +220,25 @@ class LLaMA3(BaseModel):
 			skip_special_tokens=True
 		)
 
-		def _generate() -> None:
-			assert self.model is not None
-			with torch.no_grad():
-				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)])
-				)
+		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)
+		start = time()
+
+		thread = threading.Thread(target=generate_fn)
 		thread.start()
 
 		for new_text in streamer:
-			yield new_text
+			yield new_text
+		
+		end = time()
+		total_time = end - start
+
+		self._log(f"Response generated in {total_time:.4f} seconds")