feat: 將預測任務包裝成協程

app.py

@@ -1,6 +1,7 @@
 from contextlib import asynccontextmanager
 
 import torch
+from anyio.to_thread import run_sync
 from fastapi import FastAPI
 from fastapi.params import Body
 from transformers import (
@@ -49,24 +50,30 @@ async def predict(instance: FastAPI, text: str = Body(..., embed=True)):
     :param text: 待分類的文本
     :return: 預測結果，包括文本、預測類別和置信度
     """
-    # 獲取對象
-    tokenizer, model = instance.state.tokenizer, instance.state.model
 
-    # 分詞處理
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+    def _inference():
+        # 獲取對象
+        tokenizer, model = instance.state.tokenizer, instance.state.model
 
-    # 推理
-    with torch.no_grad():
-        outputs = model(**inputs)
+        # 分詞處理
+        inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
 
-    # 處理輸出
-    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
-    confidences, classes = torch.max(predictions, dim=-1)
+        # 推理
+        with torch.no_grad():
+            outputs = model(**inputs)
+
+        # 處理輸出
+        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
+        confidences, classes = torch.max(predictions, dim=-1)
+
+        return classes.item(), confidences.item()
+
+    label, confidence = await run_sync(_inference)
 
     return {
         "text": text,
-        "label": classes.item(),
-        "confidence": confidences.item()
+        "label": label,
+        "confidence": confidence
     }