Fix model loading

Dockerfile

@@ -46,6 +46,7 @@ ENV HOME=/tmp \
     HF_DATASETS_CACHE=/tmp/.cache/huggingface/datasets \
     TRANSFORMERS_CACHE=/tmp/.cache/huggingface/transformers \
     MPLCONFIGDIR=/tmp
+RUN mkdir -p /tmp/huggingface /tmp/transformers /tmp/hub /tmp/datasets
 WORKDIR /app
 
 # 依賴

backend/utils/utils.py

@@ -17,63 +17,84 @@ 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")
+    import os
+    import torch
     
+    # 設定快取目錄
+    cache_dir = "/tmp/hf_models"
+    os.makedirs(cache_dir, exist_ok=True)
+    
+    # 設定環境變數
+    os.environ['HF_HOME'] = cache_dir
+    os.environ['TRANSFORMERS_CACHE'] = cache_dir
+    
+    # 確定設備
     if torch.cuda.is_available():
         device = torch.device("cuda")
-    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-        device = torch.device("mps")
+        dtype = torch.float16
     else:
         device = torch.device("cpu")
+        dtype = torch.float32
     
-    gpt2_aliases = {"gpt2", "openai-community/gpt2"}
+    print(f"載入模型: {model_name}")
+    print(f"設備: {device}, 精度: {dtype}")
+    print(f"快取目錄: {cache_dir}")
     
     try:
-        if model_name in gpt2_aliases:
-            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,
-                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
+        # 強制線上下載，不使用快取
+        download_kwargs = {
+            'cache_dir': cache_dir,
+            'force_download': False,  # 改為 True 如果需要強制重新下載
+            'resume_download': True,
+            'local_files_only': False
+        }
+        
+        model_kwargs = {
+            **download_kwargs,
+            'torch_dtype': dtype,
+            'low_cpu_mem_usage': True,
+        }
+        
+        if device.type == "cuda":
+            model_kwargs['device_map'] = "auto"
+        
+        # 根據模型類型載入
+        if model_name in {"gpt2", "openai-community/gpt2"}:
+            print("使用 GPT2 專用載入器")
+            tokenizer = GPT2Tokenizer.from_pretrained(model_name, **download_kwargs)
+            model = GPT2LMHeadModel.from_pretrained(model_name, **model_kwargs)
         else:
-            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,
-                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"
-            )
+            print("使用 Auto 載入器")
+            tokenizer = AutoTokenizer.from_pretrained(model_name, **download_kwargs)
+            model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
+        
+        # 設定 pad token
+        if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+            tokenizer.pad_token = tokenizer.eos_token
             
-            if getattr(model.config, "pad_token_id", None) is None and getattr(model.config, "eos_token_id", None) is not None:
+        if hasattr(model.config, 'pad_token_id') and model.config.pad_token_id is None:
+            if hasattr(model.config, 'eos_token_id') and model.config.eos_token_id is not None:
                 model.config.pad_token_id = model.config.eos_token_id
         
-        # 只在非 CPU 設備上移動模型
-        if device.type != "cpu":
-            model.to(device)
+        # 如果不是 auto device_map，手動移動到設備
+        if device.type != "cuda":
+            model = model.to(device)
             
+        print(f"✓ 成功載入模型 {model_name}")
         return tokenizer, model, device
         
+    except Exception as e:
+        import traceback
+        print(f"載入失敗: {str(e)}")
+        print(f"完整錯誤: {traceback.format_exc()}")
+        
+        # 嘗試替代方案
+        if model_name == "openai-community/gpt2":
+            print("嘗試使用 'gpt2' 替代...")
+            return load_model_and_tokenizer("gpt2")
+        
+        raise RuntimeError(f"無法載入模型 '{model_name}': {e}")
+        
     except Exception as e:
         # 提供更詳細的錯誤信息
         import traceback