Update pretrain_inference.py

pretrain_inference.py

@@ -2,6 +2,7 @@ import torch
 import torch.nn.functional as F
 from tokenizers import Tokenizer
 from novel_model import NovelTransformer
+import os
 
 # 配置参数
 VOCAB_SIZE = 8000
@@ -11,7 +12,7 @@ NUM_LAYERS = 4
 DIM_FEEDFORWARD = 512
 DROPOUT = 0.1
 MAX_LEN = 4096
-MODEL_PATH = "./novel_model/best_model.pt"
+MODEL_PATH = "./novel_model_cuda/best_model.pt"  # 修改为CUDA模型路径
 TOKENIZER_PATH = "./novel_tokenizer.json"
 
 def generate_text(model, tokenizer, prompt, max_length=100, temperature=0.8, top_k=50, top_p=0.9, device="cuda"):
@@ -65,6 +66,26 @@ def generate_text(model, tokenizer, prompt, max_length=100, temperature=0.8, top
     output = tokenizer.decode(input_ids[0].tolist())
     return output
 
+def load_model_safely(model_path, device):
+    """安全加载模型，处理设备不兼容的情况"""
+    try:
+        print(f"尝试加载模型: {model_path}")
+        # 使用map_location参数确保模型加载到正确的设备
+        checkpoint = torch.load(model_path, map_location=device)
+        return checkpoint
+    except (AssertionError, RuntimeError) as e:
+        print(f"加载模型时出错: {e}")
+        print("这可能是因为模型是在不同的设备上保存的（如Intel HPU）")
+        
+        # 尝试使用CPU作为中间设备
+        try:
+            print("尝试通过CPU加载模型...")
+            checkpoint = torch.load(model_path, map_location="cpu")
+            return checkpoint
+        except Exception as e2:
+            print(f"通过CPU加载失败: {e2}")
+            return None
+
 def main():
     # 设置设备
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -79,9 +100,25 @@ def main():
     # 加载分词器
     tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
     
-    # 加载模型
-    checkpoint = torch.load(MODEL_PATH, map_location=device)
+    # 尝试加载模型
+    model_loaded = False
+    
+    # 首先尝试加载CUDA模型
+    if os.path.exists(MODEL_PATH):
+        checkpoint = load_model_safely(MODEL_PATH, device)
+        if checkpoint is not None:
+            model_loaded = True
     
+    # 如果CUDA模型加载失败，尝试加载原始模型
+    if not model_loaded:
+        original_model_path = "d:/图像/novel_model/best_model.pt"
+        if os.path.exists(original_model_path):
+            print(f"尝试加载原始模型: {original_model_path}")
+            checkpoint = load_model_safely(original_model_path, device)
+            if checkpoint is not None:
+                model_loaded = True
+    
+    # 创建模型
     model = NovelTransformer(
         vocab_size=VOCAB_SIZE,
         d_model=D_MODEL,
@@ -92,7 +129,17 @@ def main():
         max_len=MAX_LEN
     )
     
-    model.load_state_dict(checkpoint['model_state_dict'])
+    # 如果成功加载了模型，加载状态
+    if model_loaded and checkpoint is not None:
+        try:
+            model.load_state_dict(checkpoint['model_state_dict'])
+            print("模型状态加载成功")
+        except Exception as e:
+            print(f"加载模型状态时出错: {e}")
+            print("将使用未训练的模型")
+    else:
+        print("无法加载任何模型，将使用未训练的模型")
+    
     model = model.to(device)
     
     # 设置为评估模式