Fix model loading

backend/utils/utils.py

@@ -17,34 +17,68 @@ from transformers import (
 )
 
 def load_model_and_tokenizer(model_name: str):
+    # 檢查可用記憶體
+    import psutil
+    available_memory = psutil.virtual_memory().available / 1024**3
+    print(f"Available memory: {available_memory:.2f} GB")
+    
     if torch.cuda.is_available():
         device = torch.device("cuda")
-    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():  # macOS Apple Silicon
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
         device = torch.device("mps")
     else:
         device = torch.device("cpu")
-
-    gpt2_aliases = {"gpt2", "openai-community/gpt2", "holistic-ai/gpt2-EMGSD"}
-
+    
+    gpt2_aliases = {"gpt2", "openai-community/gpt2"}
+    
     try:
         if model_name in gpt2_aliases:
-            tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+            tokenizer = GPT2Tokenizer.from_pretrained(
+                model_name,
+                cache_dir="/tmp/hf_cache"  # 使用臨時目錄
+            )
             if tokenizer.pad_token is None and tokenizer.eos_token is not None:
                 tokenizer.pad_token = tokenizer.eos_token
-            model = GPT2LMHeadModel.from_pretrained(model_name)
+                
+            model = GPT2LMHeadModel.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16 if device.type != "cpu" else torch.float32,
+                low_cpu_mem_usage=True,  # 關鍵：減少記憶體使用
+                cache_dir="/tmp/hf_cache"
+            )
+            
             if getattr(model.config, "pad_token_id", None) is None and getattr(model.config, "eos_token_id", None) is not None:
                 model.config.pad_token_id = model.config.eos_token_id
         else:
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_name,
+                cache_dir="/tmp/hf_cache"
+            )
             if tokenizer.pad_token is None and tokenizer.eos_token is not None:
                 tokenizer.pad_token = tokenizer.eos_token
-            model = AutoModelForCausalLM.from_pretrained(model_name)
+                
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16 if device.type != "cpu" else torch.float32,
+                low_cpu_mem_usage=True,  # 關鍵優化
+                device_map="auto" if torch.cuda.is_available() else None,
+                cache_dir="/tmp/hf_cache"
+            )
+            
             if getattr(model.config, "pad_token_id", None) is None and getattr(model.config, "eos_token_id", None) is not None:
                 model.config.pad_token_id = model.config.eos_token_id
-
-        model.to(device)
+        
+        # 只在非 CPU 設備上移動模型
+        if device.type != "cpu":
+            model.to(device)
+            
         return tokenizer, model, device
+        
     except Exception as e:
+        # 提供更詳細的錯誤信息
+        import traceback
+        print(f"Error loading model {model_name}: {str(e)}")
+        print(f"Traceback: {traceback.format_exc()}")
         raise RuntimeError(f"Failed to load model '{model_name}': {e}")
     
 def finetune(train_texts, tokenizer, model, num_epochs=20, output_dir='./data'):