From d40589e72016a79ac7dd666352cc0fd1b0100768 Mon Sep 17 00:00:00 2001 From: Pavan Kumar <66913595+ascender1729@users.noreply.github.com> Date: Sun, 18 May 2025 07:02:52 +0000 Subject: [PATCH] fix: use LiteLLM tokenizer for agent restoration; robust non-serializable property handling --- swarms/structs/agent_non_serializable.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/swarms/structs/agent_non_serializable.py b/swarms/structs/agent_non_serializable.py index f3bdec9e..1342899f 100644 --- a/swarms/structs/agent_non_serializable.py +++ b/swarms/structs/agent_non_serializable.py @@ -32,16 +32,12 @@ def restore_non_serializable_properties(agent): Restore non-serializable properties for the Agent instance after loading. This should be called after loading agent state from disk. """ - # Restore tokenizer if model_name is available + # Restore tokenizer using LiteLLM if available agent.tokenizer = None try: - if getattr(agent, "model_name", None): - try: - from transformers import AutoTokenizer - agent.tokenizer = AutoTokenizer.from_pretrained(agent.model_name) - except Exception: - agent.tokenizer = None - except ImportError: + from swarms.utils.litellm_tokenizer import count_tokens + agent.tokenizer = count_tokens # Assign the function as a tokenizer interface + except Exception: agent.tokenizer = None # Restore long_term_memory (dummy for demo, replace with real backend as needed)