wehub-resource-sync
40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
"""
|
|
This is an example for leveraging MLflow's auto tracing capabilities for Gemini.
|
|
|
|
For more information about MLflow Tracing, see: https://mlflow.org/docs/latest/llms/tracing/index.html
|
|
"""
|
|
|
|
import os
|
|
|
|
import mlflow
|
|
|
|
# Turn on auto tracing for Gemini by calling mlflow.gemini.autolog()
|
|
mlflow.gemini.autolog()
|
|
|
|
# Import the SDK and configure your API key.
|
|
from google import genai
|
|
|
|
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
|
|
|
|
# Use the generate_content method to generate responses to your prompts.
|
|
response = client.models.generate_content(
|
|
model="gemini-1.5-flash", contents="The opposite of hot is"
|
|
)
|
|
print(response.text)
|
|
|
|
# Also leverage the chat feature to conduct multi-turn interactions
|
|
chat = client.chats.create(model="gemini-1.5-flash")
|
|
response = chat.send_message("In one sentence, explain how a computer works to a young child.")
|
|
print(response.text)
|
|
response = chat.send_message("Okay, how about a more detailed explanation to a high schooler?")
|
|
print(response.text)
|
|
|
|
# Count tokens for your statement
|
|
response = client.models.count_tokens("The quick brown fox jumps over the lazy dog.")
|
|
print(response.total_tokens)
|
|
|
|
# Generate text embeddings for your content
|
|
text = "Hello world"
|
|
result = client.models.embed_content(model="text-embedding-004", contents=text)
|
|
print(result["embedding"])
|