update sth. at 2025-02-03 21:03:19

app.py

@@ -4,16 +4,35 @@ import torch
 import gradio as gr
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-from typing import List, Dict
+from pydantic import BaseModel, Field, root_validator
+from typing import List, Dict, Optional
 from functools import lru_cache
-import numpy as np
 from threading import Lock
 import uvicorn
 
 class EmbeddingRequest(BaseModel):
-    input: str
-    model: str = "jinaai/jina-embeddings-v3"
+    # 强制锁定模型参数
+    model: str = Field(
+        default="jinaai/jina-embeddings-v3",
+        description="此参数仅用于API兼容，实际模型固定为jinaai/jina-embeddings-v3",
+        frozen=True  # 禁止修改
+    )
+    # 支持三种输入字段
+    inputs: Optional[str] = Field(None, description="输入文本（兼容HuggingFace格式）")
+    input: Optional[str] = Field(None, description="输入文本（兼容OpenAI格式）")
+    prompt: Optional[str] = Field(None, description="输入文本（兼容Ollama格式）")
+
+    # 自动合并输入字段
+    @root_validator(pre=True)
+    def merge_input_fields(cls, values):
+        input_fields = ["inputs", "input", "prompt"]
+        for field in input_fields:
+            if values.get(field):
+                values["inputs"] = values[field]
+                break
+        else:
+            raise ValueError("必须提供 inputs/input/prompt 任一字段")
+        return values
 
 class EmbeddingResponse(BaseModel):
     status: str
@@ -21,7 +40,7 @@ class EmbeddingResponse(BaseModel):
 
 class EmbeddingService:
     def __init__(self):
-        self.model_name = "jinaai/jina-embeddings-v3"
+        self._true_model_name = "jinaai/jina-embeddings-v3"  # 硬编码模型名称
         self.max_length = 512
         self.device = torch.device("cpu")
         self.model = None
@@ -41,23 +60,22 @@ class EmbeddingService:
         try:
             from transformers import AutoTokenizer, AutoModel
             self.tokenizer = AutoTokenizer.from_pretrained(
-                self.model_name,
+                self._true_model_name,
                 trust_remote_code=True
             )
             self.model = AutoModel.from_pretrained(
-                self.model_name,
+                self._true_model_name,
                 trust_remote_code=True
             ).to(self.device)
             self.model.eval()
             torch.set_grad_enabled(False)
-            self.logger.info(f"模型加载成功，使用设备: {self.device}")
+            self.logger.info(f"强制加载模型: {self._true_model_name}")
         except Exception as e:
             self.logger.error(f"模型初始化失败: {str(e)}")
             raise
 
     @lru_cache(maxsize=1000)
     def get_embedding(self, text: str) -> List[float]:
-        """同步生成嵌入向量，带缓存"""
         with self.lock:
             try:
                 inputs = self.tokenizer(
@@ -85,7 +103,8 @@ app.add_middleware(
     allow_methods=["*"],
     allow_headers=["*"],
 )
-
+@app.post("/embed", response_model=EmbeddingResponse)
+@app.post("/api/embeddings", response_model=EmbeddingResponse)
 @app.post("/api/embed", response_model=EmbeddingResponse)
 @app.post("/v1/embeddings", response_model=EmbeddingResponse)
 @app.post("/generate_embeddings", response_model=EmbeddingResponse)
@@ -95,11 +114,10 @@ app.add_middleware(
 @app.post("/hf/v1/chat/completions", response_model=EmbeddingResponse)
 async def generate_embeddings(request: EmbeddingRequest):
     try:
-        # 使用run_in_executor避免事件循环问题
         embedding = await asyncio.get_running_loop().run_in_executor(
             None,
             embedding_service.get_embedding,
-            request.input
+            request.inputs  # 使用合并后的输入字段
         )
         return EmbeddingResponse(
             status="success",
@@ -112,7 +130,7 @@ async def generate_embeddings(request: EmbeddingRequest):
 async def root():
     return {
         "status": "active",
-        "model": embedding_service.model_name,
+        "true_model": embedding_service._true_model_name,
         "device": str(embedding_service.device)
     }
 
@@ -134,8 +152,11 @@ iface = gr.Interface(
     inputs=gr.Textbox(lines=3, label="输入文本"),
     outputs=gr.JSON(label="嵌入向量结果"),
     title="Jina Embeddings V3",
-    description="使用jina-embeddings-v3模型生成文本嵌入向量",
-    examples=[["这是一个测试句子。"]]
+    description="强制使用jinaai/jina-embeddings-v3模型（无视请求中的model参数）",
+    examples=[[
+        "Represent this sentence for searching relevant passages: "
+        "The sky is blue because of Rayleigh scattering"
+    ]]
 )
 
 @app.on_event("startup")
@@ -145,4 +166,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)