extract and fix get_device

llama_lora/lib/get_device.py

@@ -0,0 +1,15 @@
+import torch
+
+
+def get_device():
+    device ="cpu"
+    if torch.cuda.is_available():
+        device = "cuda"
+
+    try:
+        if torch.backends.mps.is_available():
+            device = "mps"
+    except:  # noqa: E722
+        pass
+
+    return device

llama_lora/lib/inference.py

@@ -1,6 +1,7 @@
 import torch
 import transformers
 
+from .get_device import get_device
 from .streaming_generation_utils import Iteratorize, Stream
 
 
@@ -62,16 +63,3 @@ def generate(
     decoded_output = tokenizer.decode(output, skip_special_tokens=True)
     yield decoded_output, output
     return
-
-
-def get_device():
-    if torch.cuda.is_available():
-        return "cuda"
-    else:
-        return "cpu"
-
-    try:
-        if torch.backends.mps.is_available():
-            return "mps"
-    except:  # noqa: E722
-        pass

llama_lora/models.py

@@ -8,19 +8,7 @@ from transformers import LlamaForCausalLM, LlamaTokenizer
 from peft import PeftModel
 
 from .globals import Global
-
-
-def get_device():
-    if torch.cuda.is_available():
-        return "cuda"
-    else:
-        return "cpu"
-
-    try:
-        if torch.backends.mps.is_available():
-            return "mps"
-    except:  # noqa: E722
-        pass
+from .lib.get_device import get_device
 
 
 def get_new_base_model(base_model_name):