no quantization

Former-commit-id: c0584f9b49e08c5fe7c6c2469fdd66685f3f3e44
pull/160/head
Kye 2 years ago
parent 721fde4aca
commit 14bb8c9503

@ -4,7 +4,6 @@ from swarms.orchestrator.autoscaler import AutoScaler
# worker
# from swarms.workers.worker_node import WorkerNode
from swarms.workers.worker import Worker
from swarms.workers.autobot import AutoBot
#boss
from swarms.boss.boss_node import BossNode

@ -5,7 +5,13 @@ import logging
from transformers import AutoModelForCausalLM, AutoTokenizer
class HuggingFaceLLM:
def __init__(self, model_id: str, device: str = None, max_length: int = 20, quantize: bool = False, quantization_config: dict = None):
def __init__(self,
model_id: str,
device: str = None,
max_length: int = 20,
quantize: bool = False,
quantization_config: dict = None
):
self.logger = logging.getLogger(__name__)
self.device = device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
self.model_id = model_id

Loading…
Cancel
Save