Spaces:
Running
Running
:zap: [Enhance] Default models and max_new_tokens
Browse files- apis/chat_api.py +1 -1
- 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=
|
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 |
-
|
16 |
-
|
|
|
|
|
|
|
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 =
|
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 |
):
|