代码初始化

inference/inference.py

@@ -1,7 +1,9 @@
+import asyncio
 import httpx
 from fastapi import APIRouter, Header, HTTPException
 from .apiModel import Payload
-import time
+import transformers
+import torch
 
 router = APIRouter()
 
@@ -26,6 +28,34 @@ async def fetch_model_response(payload: dict, headers: dict):
             raise HTTPException(status_code=response.status_code, detail=f"HTTP 错误: {e}")
 
 
+@router.post("/api-llama/")
+async def api_inference(
+        authorization: str = Header(...),
+        item: Payload = None):
+    print("请求：", item)
+    model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+    # 设置请求头
+    pipeline = transformers.pipeline(
+        "text-generation",
+        model=model_id,
+        model_kwargs={"torch_dtype": torch.bfloat16},
+        device_map="auto",
+    )
+
+    messages = [
+        {"role": "system", "content": "你是一个万能聊天机器人，能准确回答每一个提出的问题"},
+        {"role": "user", "content": "你是谁?"},
+    ]
+
+    outputs = pipeline(
+        messages,
+        max_new_tokens=256,
+    )
+
+    # 使用异步请求
+    return outputs
+
+
 @router.post("/api-inference/")
 async def api_inference(
         authorization: str = Header(...),

requirements.txt

@@ -1,4 +1,6 @@
 fastapi
 httpx
 requests
+transformers
+torch
 uvicorn[standard]