add device_map

Knowledgator · Jun 2, 2024 · f3c70c9 · f3c70c9
1 parent 363ae46
commit f3c70c9
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/unlimited_classifier/main.py b/unlimited_classifier/main.py
@@ -70,13 +70,13 @@ def initialize_model(self, model: str) -> PreTrainedModel:
         if config.is_encoder_decoder: # type: ignore
             return ( # type: ignore
                 AutoModelForSeq2SeqLM # type: ignore
-                .from_pretrained(model, device_map=self.device, quantization_config=self.quantization_config)
+                .from_pretrained(model, device_map=self.device_map, quantization_config=self.quantization_config)
             )
         else:
             try:
                 return ( # type: ignore
                     AutoModelForCausalLM # type: ignore
-                    .from_pretrained(model, device_map=self.device, quantization_config=self.quantization_config)
+                    .from_pretrained(model, device_map=self.device_map, quantization_config=self.quantization_config)
                 )
             except:
                 raise ValueError("Expected generative model.")
@@ -91,6 +91,7 @@ def __init__(
         ],
         prompt: str = "Classifity the following text:\n {}\nLabel:",
         device: str="cpu",
+        device_map:str="cpu",
         quantization_config=None,
         num_beams: int=5,
         max_new_tokens: int=512,
@@ -134,6 +135,7 @@ def __init__(
             raise ValueError("No labels provided.")
 
         self.device = device
+        self.device_map = device_map
         self.num_beams = min(num_beams, len(labels))
         self.max_new_tokens = max_new_tokens
         self.scorer = scorer