Update app.py

app.py

@@ -1,9 +1,9 @@
 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
+import spaces
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from typing import List, Dict, Any
 import time
@@ -26,20 +26,19 @@ 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]
 
+@spaces.GPU()
 def get_embedding(text: str) -> List[float]:
     inputs = tokenizer(
         text,
@@ -47,35 +46,29 @@ def get_embedding(text: str) -> List[float]:
         truncation=True,
         max_length=512,
         return_tensors="pt"
-    )
+    ).to(model.device)
     
     with torch.no_grad():
         outputs = model(**inputs)
-        embeddings = outputs.last_hidden_state[:, 0, :].numpy()
+        embeddings = outputs.last_hidden_state[:, 0, :].cpu().numpy()
     
     return embeddings[0].tolist()
 
-# OpenAI 兼容的 embeddings endpoint
 @app.post("/v1/embeddings", response_model=EmbeddingResponse)
+@spaces.GPU()
 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({
@@ -95,14 +88,10 @@ async def create_embeddings(request: EmbeddingRequest):
     
     return response
 
-# Gradio 界面
+@spaces.GPU()
 def gradio_embedding(text: str) -> Dict:
-    # 创建与 OpenAI 兼容的请求
     request = EmbeddingRequest(input=text)
-    
-    # 调用 API 处理函数
     response = create_embeddings(request)
-    
     return response.dict()
 
 # 创建 Gradio 界面
@@ -118,19 +107,9 @@ demo = gr.Interface(
     ]
 )
 
-# 启动服务
+# 挂载 Gradio 应用到 FastAPI
+app = gr.mount_gradio_app(app, demo, path="/")
+
 if __name__ == "__main__":
     import uvicorn
-    
-    # 首先启动 Gradio
-    demo.queue()
-    
-    # 然后启动 FastAPI
-    config = uvicorn.Config(
-        app=app,
-        host="0.0.0.0",
-        port=7860,
-        log_level="info"
-    )
-    server = uvicorn.Server(config)
-    server.run()
+    uvicorn.run(app, host="0.0.0.0", port=7860)