remove all explicit device assigns

utils/models.py

@@ -17,25 +17,25 @@ from .prompts import format_rag_prompt
 from .shared import generation_interrupt
 
 models = {
-    "Qwen2.5-1.5b-Instruct": "qwen/qwen2.5-1.5b-instruct",
-    "Qwen2.5-3b-Instruct": "qwen/qwen2.5-3b-instruct",
-    "Llama-3.2-1b-Instruct": "meta-llama/llama-3.2-1b-instruct",
-    "Llama-3.2-3b-Instruct": "meta-llama/llama-3.2-3b-instruct",
-    "Gemma-3-1b-it": "google/gemma-3-1b-it",
-    "Gemma-3-4b-it": "google/gemma-3-4b-it",
-    "Gemma-2-2b-it": "google/gemma-2-2b-it",
-    "Phi-4-mini-instruct": "microsoft/phi-4-mini-instruct",
-    "Cogito-v1-preview-llama-3b": "deepcogito/cogito-v1-preview-llama-3b",
-    "IBM Granite-3.3-2b-instruct": "ibm-granite/granite-3.3-2b-instruct",
-    # "Bitnet-b1.58-2B4T": "microsoft/bitnet-b1.58-2B-4T",
+    # "Qwen2.5-1.5b-Instruct": "qwen/qwen2.5-1.5b-instruct",
+    # "Qwen2.5-3b-Instruct": "qwen/qwen2.5-3b-instruct",
+    # "Llama-3.2-1b-Instruct": "meta-llama/llama-3.2-1b-instruct",
+    # "Llama-3.2-3b-Instruct": "meta-llama/llama-3.2-3b-instruct",
+    # "Gemma-3-1b-it": "google/gemma-3-1b-it",
+    # "Gemma-3-4b-it": "google/gemma-3-4b-it",
+    # "Gemma-2-2b-it": "google/gemma-2-2b-it",
+    # "Phi-4-mini-instruct": "microsoft/phi-4-mini-instruct",
+    # "Cogito-v1-preview-llama-3b": "deepcogito/cogito-v1-preview-llama-3b",
+    # "IBM Granite-3.3-2b-instruct": "ibm-granite/granite-3.3-2b-instruct",
+    "Bitnet-b1.58-2B4T": "microsoft/bitnet-b1.58-2B-4T",
     # #"MiniCPM3-RAG-LoRA": "openbmb/MiniCPM3-RAG-LoRA",
     "Qwen3-0.6b": "qwen/qwen3-0.6b",
-    "Qwen3-1.7b": "qwen/qwen3-1.7b",
-    "Qwen3-4b": "qwen/qwen3-4b",
-    "SmolLM2-1.7b-Instruct": "HuggingFaceTB/SmolLM2-1.7B-Instruct",
-    "EXAONE-3.5-2.4B-instruct": "LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct",
-    "OLMo-2-1B-Instruct": "allenai/OLMo-2-0425-1B-Instruct",
-    "icecream-3b": "aizip-dev/icecream-3b",
+    # "Qwen3-1.7b": "qwen/qwen3-1.7b",
+    # "Qwen3-4b": "qwen/qwen3-4b",
+    # "SmolLM2-1.7b-Instruct": "HuggingFaceTB/SmolLM2-1.7B-Instruct",
+    # "EXAONE-3.5-2.4B-instruct": "LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct",
+    # "OLMo-2-1B-Instruct": "allenai/OLMo-2-0425-1B-Instruct",
+    # "icecream-3b": "aizip-dev/icecream-3b",
 }
 
 tokenizer_cache = {}
@@ -159,7 +159,7 @@ def run_inference(model_name, context, question):
         if "bitnet" in model_name.lower():
             bitnet_model = BitNetForCausalLM.from_pretrained(
                 model_name,
-                device_map="auto",
+                #device_map="auto",
                 torch_dtype=torch.bfloat16,
                 #trust_remote_code=True,
             )
@@ -167,7 +167,7 @@ def run_inference(model_name, context, question):
                 "text-generation",
                 model=bitnet_model,
                 tokenizer=tokenizer,
-                device_map="auto",
+                #device_map="auto",
                 #trust_remote_code=True,
                 torch_dtype=torch.bfloat16,
                 model_kwargs={