/ examples / gateway / mosaicml / example.py
example.py
 1  from mlflow.deployments import get_deploy_client
 2  
 3  
 4  def main():
 5      client = get_deploy_client("http://localhost:7000")
 6  
 7      print(f"MosaicML endpoints: {client.list_endpoints()}\n")
 8      print(f"MosaicML completions endpoint info: {client.get_endpoint(endpoint='completions')}\n")
 9  
10      # Completions request
11      response_completions = client.predict(
12          endpoint="completions",
13          inputs={
14              "prompt": "What is the world record for flapjack consumption in a single sitting?",
15              "temperature": 0.1,
16          },
17      )
18      print(f"MosaicML response for completions: {response_completions}")
19  
20      # Embeddings request
21      response_embeddings = client.predict(
22          endpoint="embeddings",
23          inputs={"input": ["Do you carry the Storm Trooper costume in size 2T?"]},
24      )
25      print(f"MosaicML response for embeddings: {response_embeddings}")
26  
27      # Chat example
28      response_chat = client.predict(
29          endpoint="chat",
30          inputs={
31              "messages": [
32                  {
33                      "role": "system",
34                      "content": "You are a talented European rapper with a background in US history",
35                  },
36                  {
37                      "role": "user",
38                      "content": "Please recite the preamble to the US Constitution as if it were "
39                      "written today by a rapper from Reykjavík",
40                  },
41              ]
42          },
43      )
44      print(f"MosaicML response for chat: {response_chat}")
45  
46  
47  if __name__ == "__main__":
48      main()