Skip to content

Commit

Permalink
[tool cleanup]
Browse files Browse the repository at this point in the history
  • Loading branch information
kyegomez committed Jan 27, 2025
1 parent f1600d2 commit 2dc4684
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 9 deletions.
19 changes: 12 additions & 7 deletions new_features_examples/lumo_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ class Lumo:
"""
A class for generating text using the Lumo model with 4-bit quantization.
"""

def __init__(self):
"""
Initializes the Lumo model with 4-bit quantization and a tokenizer.
Expand All @@ -21,17 +22,19 @@ def __init__(self):
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
llm_int8_enable_fp32_cpu_offload=True
llm_int8_enable_fp32_cpu_offload=True,
)

self.model = LlamaForCausalLM.from_pretrained(
"lumolabs-ai/Lumo-70B-Instruct",
device_map="auto",
quantization_config=bnb_config,
use_cache=False,
attn_implementation="sdpa"
attn_implementation="sdpa",
)
self.tokenizer = AutoTokenizer.from_pretrained(
"lumolabs-ai/Lumo-70B-Instruct"
)
self.tokenizer = AutoTokenizer.from_pretrained("lumolabs-ai/Lumo-70B-Instruct")

def run(self, task: str) -> str:
"""
Expand All @@ -43,11 +46,13 @@ def run(self, task: str) -> str:
Returns:
str: The generated text.
"""
inputs = self.tokenizer(task, return_tensors="pt").to(self.model.device)
inputs = self.tokenizer(task, return_tensors="pt").to(
self.model.device
)
outputs = self.model.generate(**inputs, max_new_tokens=100)
return self.tokenizer.decode(outputs[0], skip_special_tokens=True)


return self.tokenizer.decode(
outputs[0], skip_special_tokens=True
)


Agent(
Expand Down
2 changes: 0 additions & 2 deletions swarms/tools/prebuilt/__init__.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
from swarms.tools.prebuilt.math_eval import math_eval
from swarms.tools.prebuilt.code_executor import CodeExecutor

__all__ = [
"math_eval",
"CodeExecutor",
]

0 comments on commit 2dc4684

Please sign in to comment.