| """人格预测 v2 - 单模型多头""" | |
| import json, numpy as np, torch, torch.nn as nn | |
| EMBEDDING_DIM, VEC_DIM = 768, 64 | |
| TRUNK = [256, 128, 64] | |
| HEADS = { | |
| "ocean": (["开放性","尽责性","外向性","宜人性","神经质"], 10), | |
| "four": (["力量型","活泼型","完美型","和平型"], 100), | |
| "color": (["红","蓝","黄","绿"], 100), | |
| "mbti": (["E","I","S","N","T","F","J","P"], 100), | |
| "enneagram": ([f"type_{i}" for i in range(1,10)], 100), | |
| } | |
| class PersonalityModel(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| layers=[]; prev=EMBEDDING_DIM | |
| for h in TRUNK: layers+=[nn.Linear(prev,h),nn.BatchNorm1d(h),nn.ReLU(),nn.Dropout(0.2)]; prev=h | |
| self.trunk=nn.Sequential(*layers) | |
| self.heads=nn.ModuleDict({n:nn.Sequential(nn.Linear(VEC_DIM,len(d)),nn.Sigmoid()) for n,(d,_) in HEADS.items()}) | |
| def forward(self,x): | |
| v=self.trunk(x); return {n:h(v) for n,h in self.heads.items()}, v | |
| class PersonalityPredictor: | |
| def __init__(self, model_dir=".", device=None): | |
| self.device=device or ("cuda"if torch.cuda.is_available()else"cpu") | |
| from sentence_transformers import SentenceTransformer | |
| self.enc=SentenceTransformer("shibing624/text2vec-base-chinese",device=self.device) | |
| self.model=PersonalityModel().to(self.device) | |
| ckpt=torch.load(f"{model_dir}/personality_model.pt",map_location=self.device) | |
| self.model.load_state_dict(ckpt["state_dict"]); self.model.eval() | |
| def predict(self, text): | |
| emb=self.enc.encode([text],normalize_embeddings=True) | |
| x=torch.tensor(emb,dtype=torch.float32).to(self.device) | |
| with torch.no_grad(): out,vec=self.model(x) | |
| result={} | |
| for n,(dims,scale) in HEADS.items(): | |
| s=out[n].cpu().numpy()[0]; r={} | |
| for i,d in enumerate(dims): r[d]=round(float(s[i])*scale,1) | |
| if n in("four","color"): | |
| r["_primary"]=dims[int(np.argmax(s))] | |
| elif n=="mbti": | |
| e,i_,s_,n_,t,f,j,p=s | |
| r["_mbti"]=("E"if e>i_ else"I")+("S"if s_>n_ else"N")+("T"if t>f else"F")+("J"if j>p else"P") | |
| elif n=="enneagram": | |
| idx=int(np.argmax(s)); names=["完美型","助人型","成就型","自我型","理智型","疑惑型","活跃型","领袖型","和平型"] | |
| r["_primary"]=f"{idx+1}号·{names[idx]}" | |
| result[n]=r | |
| result["vector"]=vec.cpu().numpy()[0].tolist() | |
| return result | |
| if __name__=="__main__": | |
| p=PersonalityPredictor() | |
| for t in ["曹操性格奸诈多疑、雄才大略，善于用人却心狠手辣", | |
| "此人极其暴躁冲动，一言不合就动手打人，毫无耐心"]: | |
| r=p.predict(t); print(f"\n{t[:30]}...") | |
| print(f" OCEAN:{r['ocean']} 四型:{r['four'].get('_primary')} MBTI:{r['mbti'].get('_mbti')}") | |