feat: Add device mapping for llmlingua2

MagnusS0 · Mar 25, 2024 · a0a3c43 · a0a3c43
1 parent c32ee0b
commit a0a3c43
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/huginn_hears/main.py b/huginn_hears/main.py
@@ -113,7 +113,8 @@ def load_extractive_summarizer(self):
         """
         model = PromptCompressor(
             model_name=self.model_name,
-            use_llmlingua2=True
+            use_llmlingua2=True,
+            device_map="cuda" if torch.cuda.is_available() else "cpu"
         )
         try:
             yield model