Hansimov commited on
Commit
6aa8b86
1 Parent(s): 3dfcb72

:zap: [Enhance] Default models and max_new_tokens

Browse files
Files changed (2) hide show
  1. apis/chat_api.py +1 -1
  2. networks/message_streamer.py +7 -3
apis/chat_api.py CHANGED
@@ -43,7 +43,7 @@ class ChatAPIApp:
43
  description="(float) Temperature",
44
  )
45
  max_tokens: int = Field(
46
- default=32000,
47
  description="(int) Max tokens",
48
  )
49
  stream: bool = Field(
 
43
  description="(float) Temperature",
44
  )
45
  max_tokens: int = Field(
46
+ default=8192,
47
  description="(int) Max tokens",
48
  )
49
  stream: bool = Field(
networks/message_streamer.py CHANGED
@@ -9,11 +9,15 @@ from utils.enver import enver
9
  class MessageStreamer:
10
  MODEL_MAP = {
11
  "mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1",
 
12
  }
13
 
14
  def __init__(self, model: str):
15
- self.model = model
16
- self.model_fullname = self.MODEL_MAP[model]
 
 
 
17
 
18
  def parse_line(self, line):
19
  line = line.decode("utf-8")
@@ -26,7 +30,7 @@ class MessageStreamer:
26
  self,
27
  prompt: str = None,
28
  temperature: float = 0.01,
29
- max_new_tokens: int = 32000,
30
  stream: bool = True,
31
  yield_output: bool = False,
32
  ):
 
9
  class MessageStreamer:
10
  MODEL_MAP = {
11
  "mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1",
12
+ "default": "mistralai/Mixtral-8x7B-Instruct-v0.1",
13
  }
14
 
15
  def __init__(self, model: str):
16
+ if model in self.MODEL_MAP.keys():
17
+ self.model = model
18
+ else:
19
+ self.model = "default"
20
+ self.model_fullname = self.MODEL_MAP[self.model]
21
 
22
  def parse_line(self, line):
23
  line = line.decode("utf-8")
 
30
  self,
31
  prompt: str = None,
32
  temperature: float = 0.01,
33
+ max_new_tokens: int = 8192,
34
  stream: bool = True,
35
  yield_output: bool = False,
36
  ):