Add streaming toggle

README.md

@@ -7,7 +7,7 @@ sdk: gradio
 sdk_version: 4.24.0
 app_file: app.py
 pinned: false
-short_description: 'Join the network at: https://replicantzk.com.'
+short_description: 'Learn more at https://replicantzk.com.'
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference.

app.py

@@ -1,13 +1,22 @@
+import os
 from openai import AsyncOpenAI
 import gradio as gr
 
-client = AsyncOpenAI()
+default_model = "llama3:8b-instruct-q4_K_M"
+models = ["llama3:8b-instruct-q4_K_M", "codestral:22b-v0.1-q4_K_M"]
+description = "Learn more at https://replicantzk.com."
+base_url = os.getenv("OPENAI_BASE_URL") or "https://platform.replicantzk.com"
+api_key = os.getenv("OPENAI_API_KEY")
+
+
+async def predict(message, history, model, temperature, stream, base_url, api_key):
+    client = AsyncOpenAI(base_url=base_url, api_key=api_key)
 
-async def predict(message, history, model, temperature):
     history_openai_format = []
     for human, assistant in history:
         history_openai_format.append({"role": "user", "content": human})
         history_openai_format.append({"role": "assistant", "content": assistant})
+
     history_openai_format.append({"role": "user", "content": message})
 
     try:
@@ -15,31 +24,32 @@ async def predict(message, history, model, temperature):
             model=model,
             messages=history_openai_format,
             temperature=temperature,
-            stream=True,
+            stream=stream,
         )
 
-        partial_message = ""
-        async for chunk in response:
-            if chunk.choices[0].delta.content is not None:
-                partial_message += chunk.choices[0].delta.content
-                yield partial_message
+        if stream:
+            partial_message = ""
+            async for chunk in response:
+                if chunk.choices[0].delta.content is not None:
+                    partial_message += chunk.choices[0].delta.content
+                    yield partial_message
+        else:
+            yield response.choices[0].message.content
 
     except Exception as e:
-         raise gr.Error(str(e))
-
-model_dropdown = gr.Dropdown(
-    label="Model",
-    choices=[
-        "llama3:8b-instruct-q4_K_M",
-        "codestral:22b-v0.1-q4_K_M",
-    ],
-    value="llama3:8b-instruct-q4_K_M",
+        raise gr.Error(str(e))
+
+
+model = gr.Dropdown(label="Model", choices=models, value=default_model)
+temperature = gr.Slider(0, 1, value=0, label="Temperature")
+stream = gr.Checkbox(value=True, label="Stream")
+base_url = gr.Textbox(label="OpenAI-compatible base URL", value=base_url)
+api_key = gr.Textbox(label="OpenAI-compatible API key", type="password", value=api_key)
+demo = gr.ChatInterface(
+    fn=predict,
+    additional_inputs=[model, temperature, stream, base_url, api_key],
+    description=description,
 )
-temperature_slider = gr.Slider(0, 1, value=0, label="Temperature")
-api_key = gr.Textbox(label="OpenAI API Key")
-
-desc = "Learn more and join here: https://replicantzk.com."
 
-gr.ChatInterface(
-    predict, additional_inputs=[model_dropdown, temperature_slider], description=desc
-).launch()
+if __name__ == "__main__":
+    demo.launch()

sh/reqs.sh

@@ -1,3 +1,3 @@
 #! /bin/sh
 
-poetry export -f requirements.txt --output requirements.txt
+poetry export -f requirements.txt --output ./requirements.txt

sh/start.sh

@@ -1,3 +1,3 @@
 #! /bin/sh
 
-poetry run python ./app.py
+poetry run gradio ./app.py