|
|
@@ -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)
|