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

app.py

@@ -8,7 +8,6 @@ from pydantic import BaseModel
 from typing import List, Dict
 from functools import lru_cache
 import uvicorn
-import psutil
 import numpy as np
 
 class EmbeddingRequest(BaseModel):
@@ -23,15 +22,12 @@ class EmbeddingService:
     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.setup_logging()
-        
-        # CPU优化配置
-        torch.set_num_threads(self.num_threads)
+        # CPU优化
+        torch.set_num_threads(4)
 
     def setup_logging(self):
         logging.basicConfig(
@@ -49,46 +45,49 @@ class EmbeddingService:
             )
             self.model = AutoModel.from_pretrained(
                 self.model_name,
-                trust_remote_code=True,
-                torch_dtype=torch.float32  # CPU使用float32
+                trust_remote_code=True
             ).to(self.device)
-            
             self.model.eval()
             torch.set_grad_enabled(False)
-            self.logger.info(f"模型加载成功，CPU线程数: {self.num_threads}")
+            self.logger.info(f"模型加载成功，使用设备: {self.device}")
         except Exception as e:
             self.logger.error(f"模型初始化失败: {str(e)}")
             raise
 
+    async def _generate_embedding_internal(self, text: str) -> List[float]:
+        """内部嵌入生成函数"""
+        if not text.strip():
+            raise ValueError("输入文本不能为空")
+
+        inputs = self.tokenizer(
+            text,
+            return_tensors="pt",
+            truncation=True,
+            max_length=self.max_length,
+            padding=True
+        )
+
+        with torch.no_grad():
+            outputs = self.model(**inputs).last_hidden_state.mean(dim=1)
+            return outputs.numpy().tolist()[0]
+
     @lru_cache(maxsize=1000)
-    async def generate_embedding(self, text: str) -> List[float]:
+    def get_cached_embedding(self, text: str) -> List[float]:
+        """缓存包装函数"""
+        loop = asyncio.new_event_loop()
         try:
-            inputs = self.tokenizer(
-                text,
-                return_tensors="pt",
-                truncation=True,
-                max_length=self.max_length,
-                padding=True
-            )
+            return loop.run_until_complete(self._generate_embedding_internal(text))
+        finally:
+            loop.close()
 
-            with torch.no_grad():
-                outputs = self.model(**inputs).last_hidden_state.mean(dim=1)
-                return outputs.numpy().tolist()[0]
-        except Exception as e:
-            self.logger.error(f"生成嵌入向量失败: {str(e)}")
-            raise
-
-# FastAPI应用初始化
+# 初始化服务
+embedding_service = EmbeddingService()
 app = FastAPI(
     title="Jina Embeddings API",
     description="Text embedding generation service using jina-embeddings-v3",
     version="1.0.0"
 )
 
-# 初始化服务
-embedding_service = EmbeddingService()
-
-# CORS配置
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -97,7 +96,6 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# API端点
 @app.post("/generate_embeddings", response_model=EmbeddingResponse)
 @app.post("/api/v1/embeddings", response_model=EmbeddingResponse)
 @app.post("/hf/v1/embeddings", response_model=EmbeddingResponse)
@@ -105,11 +103,13 @@ app.add_middleware(
 @app.post("/hf/v1/chat/completions", response_model=EmbeddingResponse)
 async def generate_embeddings(request: EmbeddingRequest):
     try:
-        embedding = await embedding_service.generate_embedding(request.input)
+        embedding = embedding_service.get_cached_embedding(request.input)
         return EmbeddingResponse(
             status="success",
             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))
 
@@ -118,14 +118,13 @@ async def root():
     return {
         "status": "active",
         "model": embedding_service.model_name,
-        "device": str(embedding_service.device),
-        "cpu_threads": embedding_service.num_threads
+        "device": str(embedding_service.device)
     }
 
 # Gradio界面
 def gradio_interface(text: str) -> Dict:
     try:
-        embedding = asyncio.run(embedding_service.generate_embedding(text))
+        embedding = embedding_service.get_cached_embedding(text)
         return {
             "status": "success",
             "embeddings": [embedding]
@@ -150,17 +149,11 @@ async def startup_event():
     await embedding_service.initialize()
 
 if __name__ == "__main__":
-    # 初始化服务
     asyncio.run(embedding_service.initialize())
-    
-    # 挂载Gradio应用
     gr.mount_gradio_app(app, iface, path="/ui")
-    
-    # 启动服务
     uvicorn.run(
         app,
         host="0.0.0.0",
         port=7860,
-        workers=1,
-        loop="asyncio"
+        workers=1
     )