:gem: [Feature] Enable EmbeddingApp: /encode api
Browse files- app.py +113 -0
- configs/info.json +6 -0
app.py
ADDED
@@ -0,0 +1,113 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import argparse
|
2 |
+
import markdown2
|
3 |
+
import os
|
4 |
+
import sys
|
5 |
+
import uvicorn
|
6 |
+
|
7 |
+
from pathlib import Path
|
8 |
+
from typing import Union, Optional
|
9 |
+
|
10 |
+
from fastapi import FastAPI
|
11 |
+
from pydantic import BaseModel, Field
|
12 |
+
from fastapi.responses import HTMLResponse
|
13 |
+
from tclogger import logger, OSEnver
|
14 |
+
|
15 |
+
from transforms.embed import JinaAIEmbedder
|
16 |
+
from configs.constants import AVAILABLE_MODELS
|
17 |
+
|
18 |
+
info_path = Path(__file__).parent / "configs" / "info.json"
|
19 |
+
ENVER = OSEnver(info_path)
|
20 |
+
|
21 |
+
|
22 |
+
class EmbeddingApp:
|
23 |
+
def __init__(self):
|
24 |
+
self.app = FastAPI(
|
25 |
+
docs_url="/",
|
26 |
+
title=ENVER["app_name"],
|
27 |
+
swagger_ui_parameters={"defaultModelsExpandDepth": -1},
|
28 |
+
version=ENVER["version"],
|
29 |
+
)
|
30 |
+
self.embedder = JinaAIEmbedder()
|
31 |
+
self.setup_routes()
|
32 |
+
|
33 |
+
def get_available_models(self):
|
34 |
+
return AVAILABLE_MODELS
|
35 |
+
|
36 |
+
def get_readme(self):
|
37 |
+
readme_path = Path(__file__).parents[1] / "README.md"
|
38 |
+
with open(readme_path, "r", encoding="utf-8") as rf:
|
39 |
+
readme_str = rf.read()
|
40 |
+
readme_html = markdown2.markdown(
|
41 |
+
readme_str, extras=["table", "fenced-code-blocks", "highlightjs-lang"]
|
42 |
+
)
|
43 |
+
return readme_html
|
44 |
+
|
45 |
+
class EncodePostItem(BaseModel):
|
46 |
+
text: Union[str, list[str]] = Field(
|
47 |
+
default=None,
|
48 |
+
summary="Input text(s) to embed",
|
49 |
+
)
|
50 |
+
model: Optional[str] = Field(
|
51 |
+
default=AVAILABLE_MODELS[0],
|
52 |
+
summary="Embedding model name",
|
53 |
+
)
|
54 |
+
|
55 |
+
def encode(self, item: EncodePostItem):
|
56 |
+
logger.note(f"> Encoding text: [{item.text}]", end=" ")
|
57 |
+
if item.model != self.embedder.model:
|
58 |
+
self.embedder.switch_model(item.model)
|
59 |
+
embeddings = self.embedder.encode(item.text).tolist()
|
60 |
+
logger.success(f"[{len(embeddings[0])}]")
|
61 |
+
if len(embeddings) == 1:
|
62 |
+
return embeddings[0]
|
63 |
+
else:
|
64 |
+
return embeddings
|
65 |
+
|
66 |
+
def setup_routes(self):
|
67 |
+
self.app.get(
|
68 |
+
"/models",
|
69 |
+
summary="Get available models",
|
70 |
+
)(self.get_available_models)
|
71 |
+
|
72 |
+
self.app.post(
|
73 |
+
"/encode",
|
74 |
+
summary="Encode embedding for input text",
|
75 |
+
)(self.encode)
|
76 |
+
|
77 |
+
self.app.get(
|
78 |
+
"/readme",
|
79 |
+
summary="README of HF LLM API",
|
80 |
+
response_class=HTMLResponse,
|
81 |
+
include_in_schema=False,
|
82 |
+
)(self.get_readme)
|
83 |
+
|
84 |
+
|
85 |
+
class ArgParser(argparse.ArgumentParser):
|
86 |
+
def __init__(self, *args, **kwargs):
|
87 |
+
super(ArgParser, self).__init__(*args, **kwargs)
|
88 |
+
|
89 |
+
self.add_argument(
|
90 |
+
"-s",
|
91 |
+
"--server",
|
92 |
+
type=str,
|
93 |
+
default=ENVER["server"],
|
94 |
+
help=f"Server IP ({ENVER['server']}) for Embedding API",
|
95 |
+
)
|
96 |
+
self.add_argument(
|
97 |
+
"-p",
|
98 |
+
"--port",
|
99 |
+
type=int,
|
100 |
+
default=ENVER["port"],
|
101 |
+
help=f"Server Port ({ENVER['port']}) for Embedding API",
|
102 |
+
)
|
103 |
+
|
104 |
+
self.args = self.parse_args(sys.argv[1:])
|
105 |
+
|
106 |
+
|
107 |
+
app = EmbeddingApp().app
|
108 |
+
|
109 |
+
if __name__ == "__main__":
|
110 |
+
args = ArgParser().args
|
111 |
+
uvicorn.run("__main__:app", host=args.server, port=args.port)
|
112 |
+
|
113 |
+
# python -m app
|
configs/info.json
ADDED
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"version": "0.1",
|
3 |
+
"app_name": "Embedding API",
|
4 |
+
"server": "0.0.0.0",
|
5 |
+
"port": 16666
|
6 |
+
}
|