Update app.py

app.py

@@ -107,9 +107,26 @@ def download_and_merge_model(base_model_name, lora_model_name, output_dir, devic
             device_map={"": device}
         )
         
+        old_vocab_size = base_model.get_input_embeddings().weight.shape[0]
+        print(f"原始词表大小: {old_vocab_size}")
         # 加载tokenizer
         tokenizer = AutoTokenizer.from_pretrained(base_model_name)
         
+        new_vocab_size = tokenizer.vocab_size
+        print(f"调整词表大小: {old_vocab_size} -> {new_vocab_size}")
+        
+        # 保存原始权重
+        old_embeddings = base_model.get_input_embeddings().weight.data.clone()
+        old_lm_head = base_model.lm_head.weight.data.clone()
+        
+        # 调整词表大小
+        base_model.resize_token_embeddings(new_vocab_size)
+        
+        # 复制原始权重到新的张量
+        with torch.no_grad():
+            base_model.get_input_embeddings().weight.data[:new_vocab_size] = old_embeddings[:new_vocab_size]
+            base_model.lm_head.weight.data[:new_vocab_size] = old_lm_head[:new_vocab_size]
+        
         log(f"正在加载LoRA模型: {lora_model_name}")
         log("基础模型配置:" + str(base_model.config))