use custom model

app.py

@@ -1,39 +1,38 @@
+import itertools
 import gradio as gr
 import requests
 import os
 
 
-def respond(message, history, endpoint, token, temp, max_token=0):
+def respond(message, history):
+
     if len(message.strip()) == 0:
         return "ERROR the question should not be empty"
 
-    if len(token.strip()) > 0 and len(endpoint.strip()) > 0:
-        local_token = token
-        local_endpoint = endpoint
-        custom_message = "[*] "
-    else:
-        local_token = os.environ['API_TOKEN']
-        local_endpoint = os.environ['API_ENDPOINT']
-        custom_message = ""
+
+    local_token = os.environ['API_TOKEN']
+    local_endpoint = os.environ['API_ENDPOINT']
+
 
     # Add your API token to the headers
     headers = {
         'Content-Type': 'application/json',
         'Authorization': f'Bearer {local_token}'
     }
-    q = {"inputs": {"prompt": ["Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n" +
-                               message + "\n\n### Response:\n"], "max_tokens": [int(max_token)], "temperature": [int(temp)]}}
-    #print(q)
+
+    prompt = list(itertools.chain.from_iterable(history))
+    prompt.append(message)
+    q = {"inputs": [prompt]}
     try:
         response = requests.post(local_endpoint, json=q, headers=headers, timeout=100)
         response_data = response.json(
-        )["predictions"][0]["candidates"][0]["text"]
+        )["predictions"]
     except:
         response_data = "ERROR status_code:" + \
             str(response.status_code) + " response:" + response.text
 
     #print(response.json())
-    return custom_message + response_data
+    return response_data
 
 
 demo = gr.ChatInterface(
@@ -42,22 +41,13 @@ demo = gr.ChatInterface(
     textbox=gr.Textbox(placeholder="Ask me a question",
                        container=False, scale=7),
     title="Chat with a Databricks LLM serving endpoint",
-    description="This a MPT-7b model",
+    description="This a advanced model hosted on Databricks Serving",
     examples=[["Hello"], ["What is MLflow?"], ["What is Apache Spark?"]],
     cache_examples=False,
     theme="soft",
     retry_btn=None,
     undo_btn=None,
-    clear_btn="Clear",
-    additional_inputs=[
-        gr.Textbox(label="Custom Endpoint", type="text",
-                   placeholder="https://XXXXXX.cloud.databricks.com/serving-endpoints/XXXXX/invocations"),
-        gr.Textbox(label="Custom Token", type="password",
-                   placeholder="dapiXXXXXXXXXX"),
-        gr.Slider(0, 100, label="Temp", value=0),
-        gr.Slider(1, 300, label="Max token", value=75)
-    ],
-    additional_inputs_accordion_name="Settings"
+    clear_btn="Clear"
 )
 
 if __name__ == "__main__":