Hansimov commited on
Commit
74b889b
1 Parent(s): 485cae6

:gem: [Feature] Enable EmbeddingApp: /encode api

Browse files
Files changed (2) hide show
  1. app.py +113 -0
  2. 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
+ }