Hansimov commited on
Commit
e28221f
1 Parent(s): 937bc45

:gem: [Feature] New ArgParser: Specify server ip, port and running mode

Browse files
Files changed (2) hide show
  1. apis/chat_api.py +41 -1
  2. networks/message_streamer.py +0 -1
apis/chat_api.py CHANGED
@@ -1,4 +1,7 @@
 
1
  import uvicorn
 
 
2
  from fastapi import FastAPI
3
  from pydantic import BaseModel, Field
4
  from sse_starlette.sse import EventSourceResponse
@@ -76,7 +79,44 @@ class ChatAPIApp:
76
  )(self.chat_completions)
77
 
78
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
79
  app = ChatAPIApp().app
80
 
81
  if __name__ == "__main__":
82
- uvicorn.run("__main__:app", host="0.0.0.0", port=23333, reload=True)
 
 
 
 
 
 
 
 
1
+ import argparse
2
  import uvicorn
3
+ import sys
4
+
5
  from fastapi import FastAPI
6
  from pydantic import BaseModel, Field
7
  from sse_starlette.sse import EventSourceResponse
 
79
  )(self.chat_completions)
80
 
81
 
82
+ class ArgParser(argparse.ArgumentParser):
83
+ def __init__(self, *args, **kwargs):
84
+ super(ArgParser, self).__init__(*args, **kwargs)
85
+
86
+ self.add_argument(
87
+ "-s",
88
+ "--server",
89
+ type=str,
90
+ default="0.0.0.0",
91
+ help="Server IP for HF LLM Chat API",
92
+ )
93
+ self.add_argument(
94
+ "-p",
95
+ "--port",
96
+ type=int,
97
+ default=23333,
98
+ help="Server Port for HF LLM Chat API",
99
+ )
100
+
101
+ self.add_argument(
102
+ "-d",
103
+ "--dev",
104
+ default=False,
105
+ action="store_true",
106
+ help="Run in dev mode",
107
+ )
108
+
109
+ self.args = self.parse_args(sys.argv[1:])
110
+
111
+
112
  app = ChatAPIApp().app
113
 
114
  if __name__ == "__main__":
115
+ args = ArgParser().args
116
+ if args.dev:
117
+ uvicorn.run("__main__:app", host=args.server, port=args.port, reload=True)
118
+ else:
119
+ uvicorn.run("__main__:app", host=args.server, port=args.port, reload=False)
120
+
121
+ # python -m apis.chat_api # [Docker] on product mode
122
+ # python -m apis.chat_api -d # [Dev] on develop mode
networks/message_streamer.py CHANGED
@@ -4,7 +4,6 @@ import requests
4
  from messagers.message_outputer import OpenaiStreamOutputer
5
  from utils.logger import logger
6
  from utils.enver import enver
7
- from huggingface_hub import InferenceClient
8
 
9
 
10
  class MessageStreamer:
 
4
  from messagers.message_outputer import OpenaiStreamOutputer
5
  from utils.logger import logger
6
  from utils.enver import enver
 
7
 
8
 
9
  class MessageStreamer: