David Chu commited on
Commit
fbdb66b
·
unverified ·
1 Parent(s): 01a00b6

Add modal parameter that is compatible with vllm

Browse files
Files changed (1) hide show
  1. app.py +5 -3
app.py CHANGED
@@ -17,8 +17,8 @@ CHAT_CLIENTS = {
17
  ),
18
  }
19
 
 
20
  MODEL_ARGS = {
21
- "model": "default",
22
  "max_tokens": 256,
23
  "temperature": 0.3,
24
  "frequency_penalty": 0.1,
@@ -101,7 +101,9 @@ async def assistant(chat_history, system_prompt, model_name):
101
 
102
  try:
103
  stream = await client.chat.completions.create(
104
- **MODEL_ARGS, messages=chat_history
 
 
105
  )
106
 
107
  async for chunk in stream:
@@ -152,7 +154,7 @@ with gr.Blocks(title="🦎 Salamandra & Oranguten") as demo:
152
  with gr.Row():
153
  new_chat = gr.Button("New chat")
154
 
155
- response_safety = gr.Label(show_label=False)
156
 
157
  # Event Listeners:
158
  message.submit(user, inputs=[message, chatbot], outputs=[message, chatbot]).then(
 
17
  ),
18
  }
19
 
20
+ MODEL_ID = {"Salamandra": "/models/alinia/toxic-sal"}
21
  MODEL_ARGS = {
 
22
  "max_tokens": 256,
23
  "temperature": 0.3,
24
  "frequency_penalty": 0.1,
 
101
 
102
  try:
103
  stream = await client.chat.completions.create(
104
+ **MODEL_ARGS,
105
+ model=MODEL_ID.get(model_name, "default"),
106
+ messages=chat_history,
107
  )
108
 
109
  async for chunk in stream:
 
154
  with gr.Row():
155
  new_chat = gr.Button("New chat")
156
 
157
+ response_safety = gr.Label(show_label=False, show_heading=False)
158
 
159
  # Event Listeners:
160
  message.submit(user, inputs=[message, chatbot], outputs=[message, chatbot]).then(