GenZhero/classifier

1import os2import numpy as np3from fastapi import FastAPI4from transformers import AutoTokenizer5from sklearn.preprocessing import LabelEncoder6import onnxruntime as ort7 8# -----------------------------9# 配置10# -----------------------------11ONNX_PATH = os.path.join(os.getcwd(), "minilm_classifier.onnx")12if not os.path.exists(ONNX_PATH):13    raise FileNotFoundError(f"❌ 找不到 ONNX 文件: {ONNX_PATH}")14 15# -----------------------------16# 创建 ONNX Runtime session17# -----------------------------18ort_session = ort.InferenceSession(ONNX_PATH)19 20# -----------------------------21# tokenizer 和标签 encoder22# -----------------------------23tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")24labels = ["literal", "semantic"]25label_encoder = LabelEncoder()26label_encoder.fit(labels)27 28# -----------------------------29# FastAPI 实例30# -----------------------------31app = FastAPI(title="MiniLM ONNX Classifier")32 33@app.get("/")34def root():35    return {"message": "Hello! Send GET /predict?text=your_text to classify."}36 37@app.get("/predict")38def predict(text: str):39    # tokenizer 输出 numpy 数组40    inputs = tokenizer(text, return_tensors="np", padding=True, truncation=True)41    onnx_inputs = {42        "input_ids": inputs["input_ids"].astype(np.int64),43        "attention_mask": inputs["attention_mask"].astype(np.int64)44    }45 46    # ONNX 推理47    logits = ort_session.run(None, onnx_inputs)[0]48    pred_idx = np.argmax(logits, axis=1)49    pred_label = label_encoder.inverse_transform(pred_idx)[0]50 51    return {"text": text, "prediction": pred_label}52