|
|
| import os, sys, json, torch
|
| import torch.nn.functional as F
|
| from transformers import AutoTokenizer
|
| from safetensors.torch import load_file
|
|
|
|
|
| sys.path.append(os.path.join(os.path.dirname(__file__), "common"))
|
| from models import create_model_by_name
|
|
|
| def load_model(model_dir: str):
|
| cfg_path = os.path.join(model_dir, "config.json")
|
| w_path = os.path.join(model_dir, "model.safetensors")
|
| if not (os.path.exists(cfg_path) and os.path.exists(w_path)):
|
| raise FileNotFoundError("config.json หรือ model.safetensors ไม่ครบ")
|
|
|
| with open(cfg_path, "r", encoding="utf-8") as f:
|
| cfg = json.load(f)
|
|
|
| tok = AutoTokenizer.from_pretrained(cfg["base_model"])
|
| model = create_model_by_name(cfg["arch"])
|
| state = load_file(w_path)
|
| model.load_state_dict(state)
|
| model.eval()
|
| return model, tok, cfg
|
|
|
| def predict(texts, model, tok, cfg):
|
| enc = tok(texts, padding=True, truncation=True, max_length=cfg["max_len"], return_tensors="pt")
|
| with torch.no_grad():
|
| logits = model(enc["input_ids"], enc["attention_mask"])
|
| prob = F.softmax(logits, dim=1).cpu().numpy()
|
| pred = prob.argmax(1)
|
| return pred, prob
|
|
|
| if __name__ == "__main__":
|
|
|
| MODEL_DIR = sys.argv[1] if len(sys.argv) > 1 else "cnn_bilstm"
|
|
|
| model, tok, cfg = load_model(MODEL_DIR)
|
| xs = ["อาหารอร่อยมาก บริการดี", "ไม่ประทับใจเลย ช้ามาก"]
|
| y, p = predict(xs, model, tok, cfg)
|
| labels = ["negative", "positive"]
|
| for t, yy, pp in zip(xs, y, p):
|
| print(f"{t} => {labels[yy]} | prob={pp}")
|
|
|
