llmflowstack 1.2.0__tar.gz → 1.2.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llmflowstack
-Version: 1.2.0
+Version: 1.2.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.
 Author-email: Gustavo Henrique Ferreira Cruz <gustavohferreiracruz@gmail.com>
 License: MIT

{llmflowstack-1.2.0 → llmflowstack-1.2.2}/llmflowstack/decoders/GPT_OSS.py

@@ -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(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"]
 
@@ -102,6 +102,8 @@ class GPT_OSS(BaseDecoder):
 
 		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", "")
@@ -117,6 +119,9 @@ class GPT_OSS(BaseDecoder):
 		if expected_answer:
 			assistant_text += f"<|start|>assistant<|channel|>final<|message|>{expected_answer}<|return|>"
 
+		if not expected_answer and reasoning == "Off":
+			assistant_text = "<|start|>assistant<|channel|>analysis<|message|><|end|><|start|>assistant<|channel|>final<|message|>"
+
 		return (
 			f"{system_text}{developer_text}"
 			f"<|start|>user<|message|>{data["input_text"]}<|end|>"
@@ -130,7 +135,7 @@ class GPT_OSS(BaseDecoder):
 		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.")
@@ -146,7 +151,7 @@ class GPT_OSS(BaseDecoder):
 
 	def set_reasoning_level(
 		self,
-		level: Literal["Low", "Medium", "High"]
+		level: Literal["Low", "Medium", "High", "Off"]
 	) -> None:
 		self.reasoning_level = level
 
@@ -229,6 +234,8 @@ class GPT_OSS(BaseDecoder):
 				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:
@@ -268,19 +275,26 @@ class GPT_OSS(BaseDecoder):
 			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.2.0 → llmflowstack-1.2.2}/llmflowstack/decoders/Gemma.py

@@ -270,6 +270,8 @@ class Gemma3(BaseDecoder):
 			if False:
 				yield ""
 			return
+		
+		self._log(f"Processing received input...'")
 
 		if params is None:
 			params = GenerationParams(max_new_tokens=32768)
@@ -311,8 +313,15 @@ class Gemma3(BaseDecoder):
 			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.2.0 → llmflowstack-1.2.2}/llmflowstack/decoders/LLaMA3.py

@@ -1,4 +1,5 @@
 import threading
+from functools import partial
 from time import time
 from typing import Iterator, Literal, TypedDict, cast
 
@@ -187,6 +188,8 @@ class LLaMA3(BaseDecoder):
 				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:
@@ -217,20 +220,25 @@ class LLaMA3(BaseDecoder):
 			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")

{llmflowstack-1.2.0 → llmflowstack-1.2.2}/llmflowstack/decoders/LLaMA4.py

@@ -268,6 +268,8 @@ class LLaMA4(BaseDecoder):
 				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:
@@ -308,8 +310,15 @@ class LLaMA4(BaseDecoder):
 			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.2.0 → llmflowstack-1.2.2}/llmflowstack/decoders/MedGemma.py

@@ -199,6 +199,8 @@ class MedGemma(BaseDecoder):
 				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:
@@ -239,6 +241,8 @@ class MedGemma(BaseDecoder):
 			stopping_criteria=StoppingCriteriaList([StopOnToken(self.stop_token_ids)])
 		)
 
+		start = time()
+
 		thread = threading.Thread(target=generate_fn)
 		thread.start()
 
@@ -263,4 +267,9 @@ class MedGemma(BaseDecoder):
 			else:
 				if buffer.find("<unused95>") != -1:
 					is_thinking = False
-					buffer = buffer.split("<unused95>", 1)[1]
+					buffer = buffer.split("<unused95>", 1)[1]
+		
+		end = time()
+		total_time = end - start
+
+		self._log(f"Response generated in {total_time:.4f} seconds")

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

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "llmflowstack"
-version = "1.2.0"
+version = "1.2.2"
 authors = [
   { name = "Gustavo Henrique Ferreira Cruz", email = "gustavohferreiracruz@gmail.com" }
 ]