mirror of
https://github.com/neuml/txtai.git
synced 2026-10-12 02:19:52 +00:00
* enhance documentation * Address PR comment * improvement and enhancmenet * update notebook * Address PR comment
43 KiB
43 KiB
In [ ]:
%%capture
# The main point here is to follow the pattern of other notebooks.
# For more information, kindly refer to the note in the bullet right below.
# !pip install txtai==8.1.0
# !pip install txtai[pipeline]
!pip install git+https://github.com/neuml/txtai#egg=txtai[pipeline]
# to use Vertex AI
# https://github.com/BerriAI/litellm/issues/5483
# !pip install google-cloud-aiplatform==1.75.0
!pip install google-cloud-aiplatform
# to use AWS Bedrock
# https://docs.litellm.ai/docs/providers/bedrock
# !pip install boto3==1.35.88
!pip install boto3In [ ]:
import os, getpass
from txtai import LLM, Embeddings
# https://neuml.github.io/txtai/install/
# https://neuml.github.io/txtai/pipeline/text/llm/#example
# https://neuml.github.io/txtai/embeddings/configuration/vectors/#method
# https://docs.litellm.ai/docs/embedding/supported_embedding
LLM_MODEL_CONFIG = {
'GEMINI': {
'IS_ENABLED': True,
'ENV_VAR': ['GEMINI_API_KEY'],
# https://docs.litellm.ai/docs/providers/gemini
# https://ai.google.dev/gemini-api/docs/models/gemini
'LLM_MODEL_NAME': 'gemini/gemini-pro',
# https://github.com/BerriAI/litellm/tree/12c4e7e695edb07d403dd14fc768a736638bd3d1/litellm/llms/vertex_ai
# https://github.com/BerriAI/litellm/blob/e19bb55e3b4c6a858b6e364302ebbf6633a51de5/model_prices_and_context_window.json#L2625
'TEXT_EMBEDDING_PATH': 'gemini/text-embedding-004'
},
'COHERE': {
'IS_ENABLED': False,
'ENV_VAR': ['COHERE_API_KEY'],
# https://docs.litellm.ai/docs/providers/cohere
'LLM_MODEL_NAME': 'command-light',
'TEXT_EMBEDDING_PATH': 'cohere/embed-english-v3.0'
},
'AWS_BEDROCK': {
'IS_ENABLED': False,
'ENV_VAR': ['AWS_ACCESS_KEY_ID', 'AWS_SECRET_ACCESS_KEY', 'AWS_REGION_NAME'],
# https://docs.litellm.ai/docs/providers/bedrock
'LLM_MODEL_NAME': 'bedrock/amazon.titan-text-lite-v1',
'TEXT_EMBEDDING_PATH': 'amazon.titan-embed-text-v1'
},
'MISTRAL': {
'IS_ENABLED': False,
'ENV_VAR': ['MISTRAL_API_KEY'],
# https://docs.litellm.ai/docs/providers/mistral
'LLM_MODEL_NAME': 'mistral/mistral-tiny',
'TEXT_EMBEDDING_PATH': 'mistral/mistral-embed'
},
'VERTEXAI': {
'IS_ENABLED': False,
'ENV_VAR': ['GOOGLE_APPLICATION_CREDENTIALS', 'GOOGLE_VERTEX_PROJECT', 'GOOGLE_VERTEX_LOCATION'],
'ENV_VAR_SETUP': None,
# https://docs.litellm.ai/docs/providers/vertex
'LLM_MODEL_NAME': 'vertex_ai/gemini-pro',
'TEXT_EMBEDDING_PATH': 'vertex_ai/text-embedding-004'
},
'GROQ': {
'IS_ENABLED': False,
'ENV_VAR': ['GROQ_API_KEY'],
# https://docs.litellm.ai/docs/providers/groq
# https://console.groq.com/docs/models
'LLM_MODEL_NAME': 'groq/llama3-8b-8192',
# the line below has been commented out for a reason.
# to understand why, check out the Groq section right below.
# https://groq.com/retrieval-augmented-generation-with-groq-api/
# 'TEXT_EMBEDDING_PATH': 'jinaai/jina-embeddings-v2-base-en',
},
'OPENAI': {
'IS_ENABLED': False,
'ENV_VAR': ['OPENAI_API_KEY'],
# https://docs.litellm.ai/docs/providers/openai
# https://platform.openai.com/docs/models
'LLM_MODEL_NAME': 'gpt-4o-mini-2024-07-18',
'TEXT_EMBEDDING_PATH': 'text-embedding-3-small',
},
'CLAUDE': {
'IS_ENABLED': False,
# You might need the Voyage API key, depending on the embedding model you choose. Read more below!
# 'ENV_VAR': ['ANTHROPIC_API_KEY', 'VOYAGE_API_KEY'],
'ENV_VAR': ['ANTHROPIC_API_KEY'],
# https://docs.litellm.ai/docs/providers/anthropic
# https://docs.anthropic.com/en/docs/about-claude/models
'LLM_MODEL_NAME': 'anthropic/claude-3-5-haiku-20241022',
# the line below has been commented out for a reason.
# to understand why, check out the Claude Anthropic section right below.
# https://docs.anthropic.com/en/docs/build-with-claude/embeddings
# 'TEXT_EMBEDDING_PATH': 'voyage/voyage-01',
},
'VOYAGE': {
'IS_ENABLED': False,
# You might need the Anthropic API key, depending on the LLM you choose. Read more below!
# 'ENV_VAR': ['ANTHROPIC_API_KEY', 'VOYAGE_API_KEY'],
'ENV_VAR': ['VOYAGE_API_KEY'],
# the line below has been commented out for a reason.
# to understand why, check out the Voyage AI section right below.
# https://docs.litellm.ai/docs/providers/voyage
# https://docs.voyageai.com/docs/introduction
# 'LLM_MODEL_NAME': 'anthropic/claude-3-5-haiku-20241022',
#
# https://docs.voyageai.com/docs/embeddings
'TEXT_EMBEDDING_PATH': 'voyage/voyage-01',
}
}
import litellm
# # https://github.com/BerriAI/litellm/blob/11932d0576a073d83f38a418cbdf6b2d8d4ff46f/litellm/litellm_core_utils/get_llm_provider_logic.py#L322
litellm.suppress_debug_info = True
# https://docs.litellm.ai/docs/debugging/local_debugging#set-verbose
litellm.set_verbose=False
def custom_setup ():
def vertexai_custom_setup ():
# https://docs.litellm.ai/docs/embedding/supported_embedding#usage---embedding
# https://docs.litellm.ai/docs/providers/vertex
litellm.vertex_project = os.environ['GOOGLE_VERTEX_PROJECT']
litellm.vertex_location = os.environ['GOOGLE_VERTEX_LOCATION']
LLM_MODEL_CONFIG['VERTEXAI']['ENV_VAR_SETUP'] = vertexai_custom_setup
custom_setup()
LLM_MODELS = {k: v for k, v in LLM_MODEL_CONFIG.items() if v['IS_ENABLED']}
ENV_VARS = [v['ENV_VAR'] for k, v in LLM_MODELS.items()]
ENV_VARS = [x for xs in ENV_VARS for x in xs] # flatten array
print("config set!")config set!
In [ ]:
if 'MISTRAL_API_KEY' in os.environ:
del os.environ['MISTRAL_API_KEY']In [ ]:
for ENV_VARS in [v['ENV_VAR'] for k, v in LLM_MODEL_CONFIG.items()]:
for ENV_VAR in ENV_VARS:
if ENV_VAR in os.environ:
del os.environ[ENV_VAR]
print(f'unset env var: {ENV_VAR}')
unset env var: GEMINI_API_KEY unset env var: COHERE_API_KEY unset env var: AWS_ACCESS_KEY_ID unset env var: AWS_SECRET_ACCESS_KEY unset env var: AWS_REGION_NAME unset env var: MISTRAL_API_KEY unset env var: GOOGLE_APPLICATION_CREDENTIALS unset env var: GOOGLE_VERTEX_PROJECT unset env var: GOOGLE_VERTEX_LOCATION
In [ ]:
for ENV_VAR in ENV_VARS:
os.environ[ENV_VAR] = getpass.getpass(f"enter {ENV_VAR}: ") if not ENV_VAR in os.environ else os.environ[ENV_VAR]
LLM_MODELS_WITH_CUSTOM_SETUP = {k: v for k, v in LLM_MODELS.items() if 'ENV_VAR_SETUP' in v}
for LLM_MODEL, LLM_CONFIG in LLM_MODELS_WITH_CUSTOM_SETUP.items():
print(f"running custom env var setup for {LLM_MODEL}...")
LLM_CONFIG['ENV_VAR_SETUP']()
print(f"done setup for {LLM_MODEL}")
print("all env var set!")enter GEMINI_API_KEY: ·········· enter COHERE_API_KEY: ·········· enter AWS_ACCESS_KEY_ID: ·········· enter AWS_SECRET_ACCESS_KEY: ·········· enter AWS_REGION_NAME: ·········· enter MISTRAL_API_KEY: ·········· enter GOOGLE_APPLICATION_CREDENTIALS: ·········· enter GOOGLE_VERTEX_PROJECT: ·········· enter GOOGLE_VERTEX_LOCATION: ·········· enter GROQ_API_KEY: ·········· enter OPENAI_API_KEY: ·········· enter ANTHROPIC_API_KEY: ·········· enter VOYAGE_API_KEY: ·········· running custom env var setup for VERTEXAI... done setup for VERTEXAI all env var set!
In [ ]:
# A text prompt to run through the LLM pipeline
# https://neuml.github.io/txtai/pipeline/text/llm/
LLM_PROMPT_INPUT = "Where is one place you'd go in Washington, DC?"
# The embeddings dataset is versatile! It plays with lists, datasets, or even generators.
# https://neuml.github.io/txtai/embeddings/
EMBEDDING_DATA = [
"US tops 5 million confirmed virus cases",
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
"The National Park Service warns against sacrificing slower friends in a bear attack",
"Maine man wins $1M from $25 lottery ticket",
"Make huge profits without work, earn up to $100,000 a day"
]
# https://neuml.github.io/txtai/pipeline/text/llm/
def run_llm_api(prompt="", llm_model_name=None):
if llm_model_name is None:
return
else:
# A quick note: you can skip specifying the `method` argument.
# There's an autodetection logic designed to recognize it as a `litellm` model.
# ```
# llm = LLM(LLM_MODEL_NAME, method="litellm")
# ```
llm = LLM(llm_model_name)
# OR: print(llm(LLM_PROMPT_INPUT, defaultrole="user"))
print(llm([{"role": "user", "content": prompt}]))
# https://neuml.github.io/txtai/embeddings/
def run_text_embedding_api(embedding_data=None, text_embedding_path=None):
if text_embedding_path is None:
return
else:
embeddings = Embeddings(
path=text_embedding_path,
# a quick note: you can skip specifying the `method` argument - there is autodetection logic
# method="litellm"
)
# create an index for the list of text
embeddings.index(embedding_data)
print("." * 50)
print("%-20s %s" % ("Query", "Best Match"))
print("." * 50)
# run an embeddings search for each query
for query in ("feel good story", "climate change",
"public health story", "war", "wildlife", "asia",
"lucky", "dishonest junk"):
# extract uid of first result
# search result format: (uid, score)
uid = embeddings.search(query, 1)[0][0]
# print text
print("%-20s %s" % (query, embedding_data[uid]))
# Let's LOOP THROUGH each enabled LLM and embedding model.
for LLM_MODEL, LLM_CONFIG in LLM_MODELS.items():
LLM_MODEL_NAME = LLM_CONFIG['LLM_MODEL_NAME'] if 'LLM_MODEL_NAME' in LLM_CONFIG else None
TEXT_EMBEDDING_PATH = LLM_CONFIG['TEXT_EMBEDDING_PATH'] if 'TEXT_EMBEDDING_PATH' in LLM_CONFIG else None
print("-" * 50)
print(LLM_MODEL)
# https://neuml.github.io/txtai/pipeline/text/llm/
run_llm_api(
prompt=LLM_PROMPT_INPUT,
llm_model_name=LLM_MODEL_NAME
)
# https://neuml.github.io/txtai/embeddings/
run_text_embedding_api(
embedding_data=EMBEDDING_DATA,
text_embedding_path=TEXT_EMBEDDING_PATH
)
print("-" * 50)-------------------------------------------------- GEMINI The National Mall .................................................. Query Best Match .................................................. feel good story Maine man wins $1M from $25 lottery ticket climate change Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg public health story US tops 5 million confirmed virus cases war Beijing mobilises invasion craft along coast as Taiwan tensions escalate wildlife The National Park Service warns against sacrificing slower friends in a bear attack asia Beijing mobilises invasion craft along coast as Taiwan tensions escalate lucky Maine man wins $1M from $25 lottery ticket dishonest junk Make huge profits without work, earn up to $100,000 a day -------------------------------------------------- COHERE As an AI chatbot, I do not have personal preferences or the ability to travel. However, I can suggest popular tourist destinations in Washington, DC, that many people enjoy visiting: 1. The White House: The official residence and workplace of the President of the United States is a symbol of American democracy. Visitors can take a tour of the White House to learn about its history and see the beautiful architecture. 2. National Mall: This iconic open park stretches from the United States Capitol to the Lincoln Memorial, featuring various monuments and memorials. Some notable landmarks include the Washington Monument, the Lincoln Memorial Reflecting Pool, the Vietnam Veterans Memorial, and the National World War II Memorial. 3. Smithsonian Institution: The world's largest museum and research complex offers a wealth of knowledge and cultural experiences. It includes renowned museums such as the National Air and Space Museum, National Museum of Natural History, National Museum of African American History and Culture, and many more, all offering free admission. 4. United States Capitol: A visit to the Capitol allows you to explore the seat of the United States Congress and learn about the legislative process. The Capitol also features impressive architecture and art, including the iconic Capitol Dome and the National Statuary Hall Collection. 5. National Gallery of Art: Located on the National Mall, this renowned art museum houses a vast collection of paintings, sculptures, and other artworks from the Middle Ages to the present. It offers a chance to appreciate works by famous artists like Leonardo da Vinci, Rembrandt, and Claude Monet. These are just a few highlights, but Washington, DC, offers many more attractions, including historic sites, beautiful parks, vibrant neighborhoods, and cultural institutions, ensuring there's something for everyone to enjoy. .................................................. Query Best Match .................................................. feel good story Maine man wins $1M from $25 lottery ticket climate change Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg public health story US tops 5 million confirmed virus cases war Beijing mobilises invasion craft along coast as Taiwan tensions escalate wildlife The National Park Service warns against sacrificing slower friends in a bear attack asia Beijing mobilises invasion craft along coast as Taiwan tensions escalate lucky Maine man wins $1M from $25 lottery ticket dishonest junk US tops 5 million confirmed virus cases -------------------------------------------------- AWS_BEDROCK I'd recommend the Smithsonian National Air and Space Museum. It is located in Washington, DC, and is the world's largest and most visited museum complex. .................................................. Query Best Match .................................................. feel good story Maine man wins $1M from $25 lottery ticket climate change Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg public health story US tops 5 million confirmed virus cases war Beijing mobilises invasion craft along coast as Taiwan tensions escalate wildlife The National Park Service warns against sacrificing slower friends in a bear attack asia Beijing mobilises invasion craft along coast as Taiwan tensions escalate lucky Maine man wins $1M from $25 lottery ticket dishonest junk Make huge profits without work, earn up to $100,000 a day -------------------------------------------------- MISTRAL I would love to visit the Smithsonian National Museum of Natural History in Washington, DC. It's one of the most popular and largest museums in the world, with a vast array of exhibits showcasing natural history, including dinosaur fossils, mineral and gemstone collections, and marine life. I find the diversity and depth of knowledge presented in this museum fascinating. Additionally, it's free to the public, which makes it an accessible and enjoyable experience for everyone. .................................................. Query Best Match .................................................. feel good story The National Park Service warns against sacrificing slower friends in a bear attack climate change Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg public health story The National Park Service warns against sacrificing slower friends in a bear attack war The National Park Service warns against sacrificing slower friends in a bear attack wildlife The National Park Service warns against sacrificing slower friends in a bear attack asia Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg lucky Maine man wins $1M from $25 lottery ticket dishonest junk Make huge profits without work, earn up to $100,000 a day -------------------------------------------------- VERTEXAI Washington, DC, is a fascinating city with a wealth of historical and cultural attractions. If I had the opportunity to visit, I would definitely make a stop at the Smithsonian National Air and Space Museum. This museum is home to a vast collection of aircraft and spacecraft, from the Wright brothers' first plane to the Apollo 11 command module. I am particularly interested in the history of space exploration, so I would be excited to see these iconic artifacts up close. In addition to the Air and Space Museum, there are many other places in Washington, DC, that I would like to visit. The National Mall is a beautiful park that is home to many of the city's most famous monuments, including the Washington Monument, the Lincoln Memorial, and the Vietnam Veterans Memorial. I would also like to visit the White House, the Capitol Building, and the Supreme Court. These buildings are all important symbols of American democracy, and I would be honored to have the opportunity to see them in person. Washington, DC, is a city with a rich history and a bright future. I am confident that it will continue to be a major center of culture and government for many years to come. .................................................. Query Best Match .................................................. feel good story Maine man wins $1M from $25 lottery ticket climate change Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg public health story US tops 5 million confirmed virus cases war Beijing mobilises invasion craft along coast as Taiwan tensions escalate wildlife The National Park Service warns against sacrificing slower friends in a bear attack asia Beijing mobilises invasion craft along coast as Taiwan tensions escalate lucky Maine man wins $1M from $25 lottery ticket dishonest junk Make huge profits without work, earn up to $100,000 a day -------------------------------------------------- GROQ Washington, D.C. is a city with a plethora of amazing places to visit! But, if I had to pick just one place, I'd recommend the National Mall. The National Mall is a beautiful stretch of parkland that stretches from the Lincoln Memorial to the United States Capitol Building. It's home to some of the city's most iconic landmarks, including the Washington Monument, World War II Memorial, and Vietnam Veterans Memorial. You can take a leisurely stroll along the Mall, enjoy the scenic views of the city, and stop at one of the many museums or memorials along the way. The National Mall is also a great place to people-watch, with street performers and musicians adding to the lively atmosphere. Personally, I'd recommend starting at the Lincoln Memorial, where you can take in the stunning views of the Reflecting Pool and the Washington Monument. From there, you can walk to the World War II Memorial, where you can pay your respects to the millions of Americans who served in the war. Overall, the National Mall is a must-see destination in Washington, D.C. – and it's completely free! -------------------------------------------------- OPENAI One iconic place to visit in Washington, DC, is the National Mall. It's home to many famous monuments and memorials, including the Lincoln Memorial, the Washington Monument, and the Vietnam Veterans Memorial. The expansive park offers a great opportunity to explore the history and culture of the nation, and it's a beautiful spot for walking, picnicking, and enjoying the iconic views. .................................................. Query Best Match .................................................. feel good story Maine man wins $1M from $25 lottery ticket climate change Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg public health story US tops 5 million confirmed virus cases war Beijing mobilises invasion craft along coast as Taiwan tensions escalate wildlife The National Park Service warns against sacrificing slower friends in a bear attack asia Beijing mobilises invasion craft along coast as Taiwan tensions escalate lucky Maine man wins $1M from $25 lottery ticket dishonest junk Make huge profits without work, earn up to $100,000 a day -------------------------------------------------- CLAUDE One great place to visit in Washington, DC is the Smithsonian National Air and Space Museum. It's located on the National Mall and features fascinating exhibits about aviation and space exploration, including famous aircraft and spacecraft like the Wright Brothers' plane and the Apollo 11 command module. The museum is free to enter and is very popular with visitors of all ages. -------------------------------------------------- VOYAGE .................................................. Query Best Match .................................................. feel good story Maine man wins $1M from $25 lottery ticket climate change Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg public health story Maine man wins $1M from $25 lottery ticket war The National Park Service warns against sacrificing slower friends in a bear attack wildlife The National Park Service warns against sacrificing slower friends in a bear attack asia Beijing mobilises invasion craft along coast as Taiwan tensions escalate lucky Maine man wins $1M from $25 lottery ticket dishonest junk Make huge profits without work, earn up to $100,000 a day --------------------------------------------------