Bläddra i källkod

修复重构容器没有记录

lxylxy123321 1 vecka sedan
förälder
incheckning
f901da29d0
2 ändrade filer med 20 tillägg och 14 borttagningar
  1. 19 14
      backend/app/services/model_test_service.py
  2. 1 0
      backend/requirements.txt

+ 19 - 14
backend/app/services/model_test_service.py

@@ -24,20 +24,25 @@ async def test_model(model_id: str, prompt: str, max_new_tokens: int = 128, temp
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
 
-        # 优先尝试因果语言模型加载,失败则回退到通用 AutoModel
-        try:
-            model = AutoModelForCausalLM.from_pretrained(
-                model_dir,
-                torch_dtype=torch.float16,
-                device_map="auto",
-            )
-        except (KeyError, ValueError, TypeError):
-            model = AutoModel.from_pretrained(
-                model_dir,
-                torch_dtype=torch.float16,
-                device_map="auto",
-                trust_remote_code=True,
-            )
+        # 通用加载策略:尝试多种加载方式,自动兼容各种新架构
+        model = None
+        for loader_cls, kwargs in [
+            (AutoModelForCausalLM, {"trust_remote_code": True}),
+            (AutoModel, {"trust_remote_code": True}),
+        ]:
+            try:
+                model = loader_cls.from_pretrained(
+                    model_dir,
+                    torch_dtype=torch.float16,
+                    device_map="auto",
+                    **kwargs,
+                )
+                break
+            except Exception:
+                continue
+
+        if model is None:
+            return {"error": f"Unable to load model with any available loader. Model type may not be supported yet."}
         model.eval()
 
         inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

+ 1 - 0
backend/requirements.txt

@@ -24,6 +24,7 @@ scikit-learn>=1.5.0
 pillow>=10.4.0
 huggingface_hub>=0.25.0
 modelscope>=1.15.0
+addict>=2.4.0
 pandas>=2.2.0
 pyarrow>=17.0.0
 sentencepiece>=0.2.0