Switch to Databricks Model Serving

app.py

@@ -1,4 +1,6 @@
-import openai
+import numpy as np
+import pandas as pd
+import requests
 import os
 import gradio as gr
 import json
@@ -6,16 +8,22 @@ from dotenv import load_dotenv, find_dotenv
 _ = load_dotenv(find_dotenv())
 
 
-openai.api_key  = os.getenv('OPENAI_API_KEY')
+databricks_token  = os.getenv('DATABRICKS_TOKEN')
+model_uri  = "https://dbc-eb788f31-6c73.cloud.databricks.com/serving-endpoints/Mpt-7b-tester/invocations"
 
-def get_completion(prompt, model="gpt-3.5-turbo"):
-    messages = [{"role": "user", "content": prompt}]
-    response = openai.ChatCompletion.create(
-        model=model,
-        messages=messages,
-        temperature=0, # this is the degree of randomness of the model's output
-    )
-    return response.choices[0].message["content"]
+def score_model(model_uri, databricks_token, data):
+  headers = {
+    "Authorization": f"Bearer {databricks_token}",
+    "Content-Type": "application/json",
+  }
+  data_json = json.dumps({"dataframe_split": {"index": [0], "columns": ["prompt", "temperature", "max_tokens"], "data": [[data, 0.5, 500]]}})
+  response = requests.request(method='POST', headers=headers, url=model_uri, json=data_json)
+  if response.status_code != 200:
+      raise Exception(f"Request failed with status {response.status_code}, {response.text}")
+  return response.json()
+
+def get_completion(prompt):
+    return score_model(model_uri, databricks_token, prompt)
 
 def greet(input):
 	prompt = f"""