Create app.py

app.py

@@ -0,0 +1,127 @@
+import gradio as gr
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from transformers import AutoTokenizer, AutoModel
+import torch
+import numpy as np
+from pydantic import BaseModel
+from typing import List, Dict, Any
+import time
+
+# 创建 FastAPI 应用
+app = FastAPI()
+
+# 配置 CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# 加载模型和分词器
+model_name = "BAAI/bge-m3"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModel.from_pretrained(model_name)
+model.eval()
+
+# OpenAI 兼容的请求模型
+class EmbeddingRequest(BaseModel):
+    input: List[str] | str
+    model: str | None = model_name
+    encoding_format: str | None = "float"
+    user: str | None = None
+
+# OpenAI 兼容的响应模型
+class EmbeddingResponse(BaseModel):
+    object: str = "list"
+    data: List[Dict[str, Any]]
+    model: str
+    usage: Dict[str, int]
+
+def get_embedding(text: str) -> List[float]:
+    inputs = tokenizer(
+        text,
+        padding=True,
+        truncation=True,
+        max_length=512,
+        return_tensors="pt"
+    )
+    
+    with torch.no_grad():
+        outputs = model(**inputs)
+        embeddings = outputs.last_hidden_state[:, 0, :].numpy()
+    
+    return embeddings[0].tolist()
+
+# OpenAI 兼容的 embeddings endpoint
+@app.post("/v1/embeddings", response_model=EmbeddingResponse)
+async def create_embeddings(request: EmbeddingRequest):
+    start_time = time.time()
+    
+    # 处理输入
+    if isinstance(request.input, str):
+        input_texts = [request.input]
+    else:
+        input_texts = request.input
+
+    # 获取嵌入向量
+    embeddings = []
+    total_tokens = 0
+    
+    for text in input_texts:
+        # 计算 token 数量
+        tokens = tokenizer.encode(text)
+        total_tokens += len(tokens)
+        
+        # 获取嵌入向量
+        embedding = get_embedding(text)
+        
+        embeddings.append({
+            "object": "embedding",
+            "embedding": embedding,
+            "index": len(embeddings)
+        })
+
+    response = EmbeddingResponse(
+        data=embeddings,
+        model=request.model or model_name,
+        usage={
+            "prompt_tokens": total_tokens,
+            "total_tokens": total_tokens
+        }
+    )
+    
+    return response
+
+# Gradio 界面
+def gradio_embedding(text: str) -> Dict:
+    # 创建与 OpenAI 兼容的请求
+    request = EmbeddingRequest(input=text)
+    
+    # 调用 API 处理函数
+    response = create_embeddings(request)
+    
+    return response.dict()
+
+# 创建 Gradio 界面
+iface = gr.Interface(
+    fn=gradio_embedding,
+    inputs=gr.Textbox(lines=3, placeholder="输入要进行编码的文本..."),
+    outputs=gr.Json(),
+    title="BGE-M3 Embeddings (OpenAI 兼容格式)",
+    description="输入文本，获取其对应的嵌入向量，返回格式与 OpenAI API 兼容。",
+    examples=[
+        ["这是一个示例文本。"],
+        ["人工智能正在改变世界。"]
+    ]
+)
+
+# 挂载 Gradio 应用到 FastAPI
+app = gr.mount_gradio_app(app, iface, path="/")
+
+# 启动服务
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)