wehub-resource-sync
63 lines
1.8 KiB
Python
63 lines
1.8 KiB
Python
from mlflow.deployments import get_deploy_client
|
|
|
|
|
|
def main():
|
|
client = get_deploy_client("http://localhost:7000")
|
|
|
|
print(f"Gemini endpoints: {client.list_endpoints()}\n")
|
|
print(f"Gemini completions endpoint info: {client.get_endpoint(endpoint='completions')}\n")
|
|
|
|
# Chat example
|
|
response_chat = client.predict(
|
|
endpoint="chat",
|
|
inputs={
|
|
"messages": [
|
|
{
|
|
"role": "system",
|
|
"content": "You are a talented European rapper with a background in US history",
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": "Please recite the preamble to the US Constitution as if it were "
|
|
"written today by a rapper from Reykjavík",
|
|
},
|
|
],
|
|
"temperature": 0.1,
|
|
"top_p": 1,
|
|
"n": 3,
|
|
"max_tokens": 1000,
|
|
"top_k": 40,
|
|
},
|
|
)
|
|
print(f"Gemini response for chat: {response_chat}")
|
|
|
|
# Embeddings request
|
|
response_embeddings = client.predict(
|
|
endpoint="embeddings",
|
|
inputs={
|
|
"input": [
|
|
"Describe the main differences between renewable and nonrenewable energy sources."
|
|
]
|
|
},
|
|
)
|
|
print(f"Gemini response for embeddings: {response_embeddings}\n")
|
|
|
|
# Completions request
|
|
response_completions = client.predict(
|
|
endpoint="completions",
|
|
inputs={
|
|
"prompt": "Describe the main differences between renewable and nonrenewable energy sources.",
|
|
"temperature": 0.1,
|
|
"stop": ["."],
|
|
"n": 3,
|
|
"max_tokens": 100,
|
|
"top_k": 40,
|
|
"top_p": 0.5,
|
|
},
|
|
)
|
|
print(f"Gemini response for completions: {response_completions}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|