update sth. at 2025-01-16 22:21:25

app.py

@@ -1,28 +1,15 @@
 import asyncio
 import logging
-import time
 import torch
 import gradio as gr
-from fastapi import FastAPI, Request, HTTPException
+from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel, BaseSettings
-from transformers import AutoTokenizer, AutoModel
+from pydantic import BaseModel
 from typing import List, Dict
 from functools import lru_cache
-import numpy as np
 import uvicorn
-
-class Settings(BaseSettings):
-    model_name: str = "jinaai/jina-embeddings-v3"
-    max_length: int = 512
-    batch_size: int = 32
-    host: str = "0.0.0.0"
-    port: int = 7860
-    enable_gpu: bool = True
-    queue_size: int = 100
-
-    class Config:
-        env_file = ".env"
+import psutil
+import numpy as np
 
 class EmbeddingRequest(BaseModel):
     input: str
@@ -33,13 +20,18 @@ class EmbeddingResponse(BaseModel):
     embeddings: List[List[float]]
 
 class EmbeddingService:
-    def __init__(self, settings: Settings):
-        self.settings = settings
-        self.device = torch.device("cuda" if torch.cuda.is_available() and settings.enable_gpu else "cpu")
+    def __init__(self):
+        self.model_name = "jinaai/jina-embeddings-v3"
+        self.max_length = 512
+        self.batch_size = 8
+        self.device = torch.device("cpu")
+        self.num_threads = min(psutil.cpu_count(), 4)  # 限制CPU线程数
         self.model = None
         self.tokenizer = None
-        self.request_queue = asyncio.Queue(maxsize=settings.queue_size)
         self.setup_logging()
+        
+        # CPU优化配置
+        torch.set_num_threads(self.num_threads)
 
     def setup_logging(self):
         logging.basicConfig(
@@ -50,54 +42,53 @@ class EmbeddingService:
 
     async def initialize(self):
         try:
+            from transformers import AutoTokenizer, AutoModel
             self.tokenizer = AutoTokenizer.from_pretrained(
-                self.settings.model_name,
+                self.model_name,
                 trust_remote_code=True
             )
             self.model = AutoModel.from_pretrained(
-                self.settings.model_name,
-                trust_remote_code=True
+                self.model_name,
+                trust_remote_code=True,
+                torch_dtype=torch.float32  # CPU使用float32
             ).to(self.device)
+            
             self.model.eval()
-            self.logger.info(f"模型加载成功，使用设备: {self.device}")
+            torch.set_grad_enabled(False)
+            self.logger.info(f"模型加载成功，CPU线程数: {self.num_threads}")
         except Exception as e:
             self.logger.error(f"模型初始化失败: {str(e)}")
             raise
 
     @lru_cache(maxsize=1000)
-    async def generate_embedding(self, text: str) -> np.ndarray:
+    async def generate_embedding(self, text: str) -> List[float]:
         try:
             inputs = self.tokenizer(
                 text,
                 return_tensors="pt",
                 truncation=True,
-                max_length=self.settings.max_length
-            ).to(self.device)
+                max_length=self.max_length,
+                padding=True
+            )
 
             with torch.no_grad():
                 outputs = self.model(**inputs).last_hidden_state.mean(dim=1)
-                return outputs.cpu().numpy()
+                return outputs.numpy().tolist()[0]
         except Exception as e:
             self.logger.error(f"生成嵌入向量失败: {str(e)}")
             raise
 
-    async def handle_request(self, text: str) -> np.ndarray:
-        if not text.strip():
-            raise ValueError("输入文本不能为空")
-        return await self.generate_embedding(text)
-
-# 初始化服务
-settings = Settings()
-embedding_service = EmbeddingService(settings)
-
-# FastAPI应用
+# FastAPI应用初始化
 app = FastAPI(
     title="Jina Embeddings API",
     description="Text embedding generation service using jina-embeddings-v3",
     version="1.0.0"
 )
 
-# CORS中间件
+# 初始化服务
+embedding_service = EmbeddingService()
+
+# CORS配置
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -106,7 +97,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# FastAPI路由
+# API端点
 @app.post("/generate_embeddings", response_model=EmbeddingResponse)
 @app.post("/api/v1/embeddings", response_model=EmbeddingResponse)
 @app.post("/hf/v1/embeddings", response_model=EmbeddingResponse)
@@ -114,13 +105,11 @@ app.add_middleware(
 @app.post("/hf/v1/chat/completions", response_model=EmbeddingResponse)
 async def generate_embeddings(request: EmbeddingRequest):
     try:
-        embedding = await embedding_service.handle_request(request.input)
+        embedding = await embedding_service.generate_embedding(request.input)
         return EmbeddingResponse(
             status="success",
-            embeddings=embedding.tolist()
+            embeddings=[embedding]
         )
-    except ValueError as e:
-        raise HTTPException(status_code=400, detail=str(e))
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -128,18 +117,18 @@ async def generate_embeddings(request: EmbeddingRequest):
 async def root():
     return {
         "status": "active",
-        "model": settings.model_name,
-        "device": embedding_service.device,
-        "usage": "Send POST request to /generate_embeddings or use UI interface"
+        "model": embedding_service.model_name,
+        "device": str(embedding_service.device),
+        "cpu_threads": embedding_service.num_threads
     }
 
 # Gradio界面
 def gradio_interface(text: str) -> Dict:
     try:
-        embedding = asyncio.run(embedding_service.handle_request(text))
+        embedding = asyncio.run(embedding_service.generate_embedding(text))
         return {
             "status": "success",
-            "embeddings": embedding.tolist()
+            "embeddings": [embedding]
         }
     except Exception as e:
         return {
@@ -153,10 +142,7 @@ iface = gr.Interface(
     outputs=gr.JSON(label="嵌入向量结果"),
     title="Jina Embeddings V3",
     description="使用jina-embeddings-v3模型生成文本嵌入向量",
-    examples=[
-        ["这是一个测试句子。"],
-        ["人工智能正在改变世界。"]
-    ]
+    examples=[["这是一个测试句子。"]]
 )
 
 @app.on_event("startup")
@@ -164,15 +150,17 @@ async def startup_event():
     await embedding_service.initialize()
 
 if __name__ == "__main__":
-    # 确保模型初始化
+    # 初始化服务
     asyncio.run(embedding_service.initialize())
     
-    # 启动Gradio和FastAPI
+    # 挂载Gradio应用
     gr.mount_gradio_app(app, iface, path="/ui")
     
+    # 启动服务
     uvicorn.run(
         app,
-        host=settings.host,
-        port=settings.port,
-        workers=1  # GPU模式下建议使用单进程
+        host="0.0.0.0",
+        port=7860,
+        workers=1,
+        loop="asyncio"
     )