File size: 1,769 Bytes
1d6f4c0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | import re, json, os, torch, torch.nn as nn
HERE = os.path.dirname(os.path.abspath(__file__))
cfg = json.load(open(os.path.join(HERE, "vibe_config.json")))
vocab = json.load(open(os.path.join(HERE, "vibe_vocab.json")))
MAX_LEN, EMB, LAYERS, HEADS, FF = cfg["max_len"], cfg["emb"], cfg["layers"], cfg["heads"], cfg["ff"]
LABELS = cfg["labels"]
class VibeNet(nn.Module):
def __init__(self, V):
super().__init__()
self.emb = nn.Embedding(V, EMB, padding_idx=0)
self.pos = nn.Embedding(MAX_LEN, EMB)
self.drop = nn.Dropout(0.0)
layer = nn.TransformerEncoderLayer(EMB, HEADS, FF, 0.0, batch_first=True, activation="gelu")
self.enc = nn.TransformerEncoder(layer, LAYERS, enable_nested_tensor=False)
self.head = nn.Linear(EMB, 2)
def forward(self, x):
mask = (x == 0)
pos = torch.arange(x.size(1)).unsqueeze(0)
h = self.emb(x) + self.pos(pos)
h = self.enc(h, src_key_padding_mask=mask)
keep = (~mask).unsqueeze(-1).float()
pooled = (h * keep).sum(1) / keep.sum(1).clamp(min=1)
return self.head(pooled)
_model = VibeNet(len(vocab))
_model.load_state_dict(torch.load(os.path.join(HERE, "vibe_model.pt"), map_location="cpu"))
_model.eval()
_tok = lambda s: re.findall(r"[a-z']+", s.lower())
def _encode(s):
ids = [vocab.get(w, 1) for w in _tok(s)][:MAX_LEN]
if not ids: ids = [1]
return ids + [0]*(MAX_LEN-len(ids))
@torch.no_grad()
def predict_vibe(text):
text = (text or "").strip()
if not text: return {"label": "Neutral", "confidence": 0.0}
x = torch.tensor([_encode(text)])
probs = torch.softmax(_model(x), dim=1)[0]
i = int(probs.argmax())
return {"label": LABELS[i], "confidence": round(float(probs[i])*100, 1)}
|