update sth. at 2025-01-16 22:16:33

README.md

@@ -24,10 +24,11 @@ You can generate embeddings by sending a POST request to one of the following en
 
 Example request using `curl`:
 
-```sh
+``` bash
 curl -X POST https://sanbo1200-jina-embeddings-v3.hf.space/api/v1/embeddings \
   -H "Content-Type: application/json" \
   -d '{
     "input": "Your text string goes here",
     "model": "jinaai/jina-embeddings-v3"
-  }'
+  }'
+```

app.py

@@ -1,11 +1,29 @@
-from fastapi import FastAPI, HTTPException, Request
-from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModel
+import asyncio
+import logging
+import time
 import torch
+import gradio as gr
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, BaseSettings
+from transformers import AutoTokenizer, AutoModel
 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"
+
 class EmbeddingRequest(BaseModel):
     input: str
     model: str = "jinaai/jina-embeddings-v3"
@@ -14,18 +32,81 @@ class EmbeddingResponse(BaseModel):
     status: str
     embeddings: List[List[float]]
 
-# 创建FastAPI应用
+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")
+        self.model = None
+        self.tokenizer = None
+        self.request_queue = asyncio.Queue(maxsize=settings.queue_size)
+        self.setup_logging()
+
+    def setup_logging(self):
+        logging.basicConfig(
+            level=logging.INFO,
+            format='%(asctime)s - %(levelname)s - %(message)s'
+        )
+        self.logger = logging.getLogger(__name__)
+
+    async def initialize(self):
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.settings.model_name,
+                trust_remote_code=True
+            )
+            self.model = AutoModel.from_pretrained(
+                self.settings.model_name,
+                trust_remote_code=True
+            ).to(self.device)
+            self.model.eval()
+            self.logger.info(f"模型加载成功，使用设备: {self.device}")
+        except Exception as e:
+            self.logger.error(f"模型初始化失败: {str(e)}")
+            raise
+
+    @lru_cache(maxsize=1000)
+    async def generate_embedding(self, text: str) -> np.ndarray:
+        try:
+            inputs = self.tokenizer(
+                text,
+                return_tensors="pt",
+                truncation=True,
+                max_length=self.settings.max_length
+            ).to(self.device)
+
+            with torch.no_grad():
+                outputs = self.model(**inputs).last_hidden_state.mean(dim=1)
+                return outputs.cpu().numpy()
+        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应用
 app = FastAPI(
     title="Jina Embeddings API",
     description="Text embedding generation service using jina-embeddings-v3",
     version="1.0.0"
 )
 
-# 加载模型和分词器
-model_name = "jinaai/jina-embeddings-v3"
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
+# CORS中间件
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
+# FastAPI路由
 @app.post("/generate_embeddings", response_model=EmbeddingResponse)
 @app.post("/api/v1/embeddings", response_model=EmbeddingResponse)
 @app.post("/hf/v1/embeddings", response_model=EmbeddingResponse)
@@ -33,17 +114,13 @@ model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
 @app.post("/hf/v1/chat/completions", response_model=EmbeddingResponse)
 async def generate_embeddings(request: EmbeddingRequest):
     try:
-        # 使用分词器处理输入文本
-        inputs = tokenizer(request.input, return_tensors="pt", truncation=True, max_length=512)
-
-        # 生成嵌入
-        with torch.no_grad():
-            embeddings = model(**inputs).last_hidden_state.mean(dim=1)
-
+        embedding = await embedding_service.handle_request(request.input)
         return EmbeddingResponse(
             status="success",
-            embeddings=embeddings.numpy().tolist()
+            embeddings=embedding.tolist()
         )
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -51,9 +128,51 @@ async def generate_embeddings(request: EmbeddingRequest):
 async def root():
     return {
         "status": "active",
-        "model": model_name,
-        "usage": "Send POST request to /generate_embeddings or /api/v1/embeddings or /hf/v1/embeddings"
+        "model": settings.model_name,
+        "device": embedding_service.device,
+        "usage": "Send POST request to /generate_embeddings or use UI interface"
     }
 
+# Gradio界面
+def gradio_interface(text: str) -> Dict:
+    try:
+        embedding = asyncio.run(embedding_service.handle_request(text))
+        return {
+            "status": "success",
+            "embeddings": embedding.tolist()
+        }
+    except Exception as e:
+        return {
+            "status": "error",
+            "message": str(e)
+        }
+
+iface = gr.Interface(
+    fn=gradio_interface,
+    inputs=gr.Textbox(lines=3, label="输入文本"),
+    outputs=gr.JSON(label="嵌入向量结果"),
+    title="Jina Embeddings V3",
+    description="使用jina-embeddings-v3模型生成文本嵌入向量",
+    examples=[
+        ["这是一个测试句子。"],
+        ["人工智能正在改变世界。"]
+    ]
+)
+
+@app.on_event("startup")
+async def startup_event():
+    await embedding_service.initialize()
+
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    # 确保模型初始化
+    asyncio.run(embedding_service.initialize())
+    
+    # 启动Gradio和FastAPI
+    gr.mount_gradio_app(app, iface, path="/ui")
+    
+    uvicorn.run(
+        app,
+        host=settings.host,
+        port=settings.port,
+        workers=1  # GPU模式下建议使用单进程
+    )

app.py1

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModel
 import torch
@@ -6,8 +6,9 @@ from typing import List, Dict
 import uvicorn
 
 # 定义请求和响应模型
-class TextRequest(BaseModel):
-    text: str
+class EmbeddingRequest(BaseModel):
+    input: str
+    model: str = "jinaai/jina-embeddings-v3"
 
 class EmbeddingResponse(BaseModel):
     status: str
@@ -26,15 +27,19 @@ tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
 
 @app.post("/generate_embeddings", response_model=EmbeddingResponse)
-async def generate_embeddings(request: TextRequest):
+@app.post("/api/v1/embeddings", response_model=EmbeddingResponse)
+@app.post("/hf/v1/embeddings", response_model=EmbeddingResponse)
+@app.post("/api/v1/chat/completions", response_model=EmbeddingResponse)
+@app.post("/hf/v1/chat/completions", response_model=EmbeddingResponse)
+async def generate_embeddings(request: EmbeddingRequest):
     try:
         # 使用分词器处理输入文本
-        inputs = tokenizer(request.text, return_tensors="pt", truncation=True, max_length=512)
-        
+        inputs = tokenizer(request.input, return_tensors="pt", truncation=True, max_length=512)
+
         # 生成嵌入
         with torch.no_grad():
             embeddings = model(**inputs).last_hidden_state.mean(dim=1)
-        
+
         return EmbeddingResponse(
             status="success",
             embeddings=embeddings.numpy().tolist()
@@ -47,8 +52,8 @@ async def root():
     return {
         "status": "active",
         "model": model_name,
-        "usage": "Send POST request to /generate_embeddings"
+        "usage": "Send POST request to /generate_embeddings or /api/v1/embeddings or /hf/v1/embeddings"
     }
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)