You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
126 lines
4.4 KiB
126 lines
4.4 KiB
"""
|
|
Example usage of WatsonX embedding and LLM in DeepSearcher.
|
|
|
|
This example demonstrates how to configure and use IBM WatsonX
|
|
embedding models and language models with DeepSearcher.
|
|
"""
|
|
|
|
import os
|
|
from deepsearcher.configuration import Configuration
|
|
|
|
def main():
|
|
"""Example of using WatsonX with DeepSearcher."""
|
|
|
|
# Initialize configuration
|
|
config = Configuration()
|
|
|
|
# Set up environment variables (alternatively, set these in your shell)
|
|
# os.environ["WATSONX_APIKEY"] = "your-watsonx-api-key"
|
|
# os.environ["WATSONX_URL"] = "https://your-watsonx-instance.com"
|
|
# os.environ["WATSONX_PROJECT_ID"] = "your-project-id"
|
|
|
|
# Example 1: Configure WatsonX Embedding
|
|
print("=== WatsonX Embedding Configuration ===")
|
|
|
|
# Basic configuration with default model
|
|
config.set_provider_config("embedding", "WatsonXEmbedding", {})
|
|
|
|
# Configuration with custom model
|
|
config.set_provider_config("embedding", "WatsonXEmbedding", {
|
|
"model": "ibm/slate-125m-english-rtrvr-v2"
|
|
})
|
|
|
|
# Configuration with explicit credentials
|
|
# config.set_provider_config("embedding", "WatsonXEmbedding", {
|
|
# "model": "sentence-transformers/all-minilm-l6-v2",
|
|
# "api_key": "your-api-key",
|
|
# "url": "https://your-watsonx-instance.com",
|
|
# "project_id": "your-project-id"
|
|
# })
|
|
|
|
print("WatsonX Embedding configured successfully!")
|
|
|
|
# Example 2: Configure WatsonX LLM
|
|
print("\n=== WatsonX LLM Configuration ===")
|
|
|
|
# Basic configuration with default model
|
|
config.set_provider_config("llm", "WatsonX", {})
|
|
|
|
# Configuration with custom model and parameters
|
|
config.set_provider_config("llm", "WatsonX", {
|
|
"model": "ibm/granite-3-3-8b-instruct",
|
|
"max_new_tokens": 1000,
|
|
"temperature": 0.7,
|
|
"top_p": 0.9,
|
|
"top_k": 50
|
|
})
|
|
|
|
# Configuration with IBM Granite model
|
|
config.set_provider_config("llm", "WatsonX", {
|
|
"model": "ibm/granite-3-3-8b-instruct",
|
|
"max_new_tokens": 512,
|
|
"temperature": 0.1
|
|
})
|
|
|
|
print("WatsonX LLM configured successfully!")
|
|
|
|
# Example 3: Test embedding functionality
|
|
print("\n=== Testing WatsonX Embedding ===")
|
|
try:
|
|
from deepsearcher.embedding.watsonx_embedding import WatsonXEmbedding
|
|
|
|
# Check if environment variables are set
|
|
if all(os.getenv(var) for var in ["WATSONX_APIKEY", "WATSONX_URL", "WATSONX_PROJECT_ID"]):
|
|
embedding = WatsonXEmbedding()
|
|
|
|
# Test single query embedding
|
|
query = "What is artificial intelligence?"
|
|
query_embedding = embedding.embed_query(query)
|
|
print(f"Query embedding dimension: {len(query_embedding)}")
|
|
|
|
# Test document embeddings
|
|
documents = [
|
|
"Artificial intelligence is a branch of computer science.",
|
|
"Machine learning is a subset of AI.",
|
|
"Deep learning uses neural networks."
|
|
]
|
|
doc_embeddings = embedding.embed_documents(documents)
|
|
print(f"Document embeddings: {len(doc_embeddings)} vectors of dimension {len(doc_embeddings[0])}")
|
|
|
|
else:
|
|
print("Environment variables not set. Skipping embedding test.")
|
|
|
|
except ImportError:
|
|
print("WatsonX dependencies not installed. Run: pip install ibm-watsonx-ai")
|
|
except Exception as e:
|
|
print(f"Error testing embedding: {e}")
|
|
|
|
# Example 4: Test LLM functionality
|
|
print("\n=== Testing WatsonX LLM ===")
|
|
try:
|
|
from deepsearcher.llm.watsonx import WatsonX
|
|
|
|
# Check if environment variables are set
|
|
if all(os.getenv(var) for var in ["WATSONX_APIKEY", "WATSONX_URL", "WATSONX_PROJECT_ID"]):
|
|
llm = WatsonX()
|
|
|
|
# Test chat functionality
|
|
messages = [
|
|
{"role": "system", "content": "You are a helpful AI assistant."},
|
|
{"role": "user", "content": "Explain what artificial intelligence is in one sentence."}
|
|
]
|
|
|
|
response = llm.chat(messages)
|
|
print(f"LLM Response: {response.content}")
|
|
print(f"Tokens used: {response.total_tokens}")
|
|
|
|
else:
|
|
print("Environment variables not set. Skipping LLM test.")
|
|
|
|
except ImportError:
|
|
print("WatsonX dependencies not installed. Run: pip install ibm-watsonx-ai")
|
|
except Exception as e:
|
|
print(f"Error testing LLM: {e}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|