update sth. at 2025-01-16 22:31:26

README.md

@@ -31,4 +31,5 @@ curl -X POST https://sanbo1200-jina-embeddings-v3.hf.space/api/v1/embeddings \
     "input": "Your text string goes here",
     "model": "jinaai/jina-embeddings-v3"
   }'
-```
+```
+

app.py

@@ -7,8 +7,8 @@ from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from typing import List, Dict
 from functools import lru_cache
-import uvicorn
 import numpy as np
+from threading import Lock
 
 class EmbeddingRequest(BaseModel):
     input: str
@@ -25,9 +25,9 @@ class EmbeddingService:
         self.device = torch.device("cpu")
         self.model = None
         self.tokenizer = None
+        self.lock = Lock()
         self.setup_logging()
-        # CPU优化
-        torch.set_num_threads(4)
+        torch.set_num_threads(4)  # CPU优化
 
     def setup_logging(self):
         logging.basicConfig(
@@ -54,39 +54,28 @@ class EmbeddingService:
             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)
-    def get_cached_embedding(self, text: str) -> List[float]:
-        """缓存包装函数"""
-        loop = asyncio.new_event_loop()
-        try:
-            return loop.run_until_complete(self._generate_embedding_internal(text))
-        finally:
-            loop.close()
+    def get_embedding(self, text: str) -> List[float]:
+        """同步生成嵌入向量，带缓存"""
+        with self.lock:
+            try:
+                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]
+            except Exception as e:
+                self.logger.error(f"生成嵌入向量失败: {str(e)}")
+                raise
 
-# 初始化服务
 embedding_service = EmbeddingService()
-app = FastAPI(
-    title="Jina Embeddings API",
-    description="Text embedding generation service using jina-embeddings-v3",
-    version="1.0.0"
-)
+app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
@@ -103,13 +92,16 @@ app.add_middleware(
 @app.post("/hf/v1/chat/completions", response_model=EmbeddingResponse)
 async def generate_embeddings(request: EmbeddingRequest):
     try:
-        embedding = embedding_service.get_cached_embedding(request.input)
+        # 使用run_in_executor避免事件循环问题
+        embedding = await asyncio.get_running_loop().run_in_executor(
+            None,
+            embedding_service.get_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))
 
@@ -121,10 +113,9 @@ async def root():
         "device": str(embedding_service.device)
     }
 
-# Gradio界面
 def gradio_interface(text: str) -> Dict:
     try:
-        embedding = embedding_service.get_cached_embedding(text)
+        embedding = embedding_service.get_embedding(text)
         return {
             "status": "success",
             "embeddings": [embedding]
@@ -151,9 +142,4 @@ async def startup_event():
 if __name__ == "__main__":
     asyncio.run(embedding_service.initialize())
     gr.mount_gradio_app(app, iface, path="/ui")
-    uvicorn.run(
-        app,
-        host="0.0.0.0",
-        port=7860,
-        workers=1
-    )
+    uvicorn.run(app, host="0.0.0.0", port=7860, workers=1)

requirements.txt

@@ -3,4 +3,9 @@ torch
 einops
 fastapi
 uvicorn
-pydantic
+pydantic
+gradio
+numpy
+python-multipart
+sentencepiece
+safetensors