集成：mixedbread ai

在几秒钟内使用 mixedbread 的模型以及顶级的开源模型

作者

mixedbread ai

GitHub Repo PyPI Package

概述

mixedbread ai 是一家人工智能初创公司，提供开源和内部的嵌入（embedding）和重排序（reranking）模型。您可以从各种基础模型中进行选择，找到最适合您用例的模型。更多信息可在文档页面找到。

安装

使用简单的 pip 命令安装 mixedbread ai 集成

pip install mixedbread-ai-haystack

使用

此集成包含 3 个组件

对于文档，您可以使用 MixedbreadAIDocumentEmbedder；对于查询，您可以使用 MixedbreadAITextEmbedder。一旦您为您的特定用例选择了组件，请使用 model 和 api_key 来初始化组件。您也可以设置环境变量 MXBAI_API_KEY 来代替将 api key 作为参数传递。

管道中的 Embedders

from haystack import Document, Pipeline
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.components.writers import DocumentWriter
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from mixedbread_ai_haystack.embedders import MixedbreadAIDocumentEmbedder, MixedbreadAITextEmbedder

# Set-up the Document Store and Documents
document_store = InMemoryDocumentStore(embedding_similarity_function="cosine")
documents = [
    Document(content="china is the most populous country in the world."), 
    Document(content="india is the second most populous country in the world."), 
    Document(content="united states is the third most populous country in the world.")
]

# Indexing Pipeline
indexing_pipeline = Pipeline()
indexing_pipeline.add_component("doc_embedder", MixedbreadAIDocumentEmbedder(model="mixedbread-ai/mxbai-embed-large-v1"))
indexing_pipeline.add_component("writer", DocumentWriter(document_store=document_store))
indexing_pipeline.connect("doc_embedder", "writer")

indexing_pipeline.run({"doc_embedder": {"documents": documents}})

# Query Pipeline
query_pipeline = Pipeline()
query_pipeline.add_component("text_embedder",  MixedbreadAITextEmbedder(model="mixedbread-ai/mxbai-embed-large-v1"))
query_pipeline.add_component("retriever", InMemoryEmbeddingRetriever(document_store=document_store))
query_pipeline.connect("text_embedder.embedding", "retriever.query_embedding")

results = query_pipeline.run({"text_embedder": {"text": "Which country has the biggest population?"}})
top_document = results["retriever"]["documents"][0].content
print(top_document)

管道中的 Reranker

from haystack import Document, Pipeline
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from mixedbread_ai_haystack.rerankers import MixedbreadAIReranker

# Set-up the Document Store and Documents
documents = [
    Document(content="china is the most populous country in the world."),
    Document(content="india is the second most populous country in the world."),
    Document(content="united states is the third most populous country in the world.")
]
document_store = InMemoryDocumentStore()
document_store.write_documents(documents)

# Define the Retriever and Reranker
retriever = InMemoryBM25Retriever(document_store=document_store)
reranker = MixedbreadAIReranker(model="mixedbread-ai/mxbai-rerank-large-v1", top_k=3)

# Rerank Pipeline
rerank_pipeline = Pipeline()
rerank_pipeline.add_component("retriever", retriever)
rerank_pipeline.add_component("reranker", reranker)
rerank_pipeline.connect("retriever.documents", "reranker.documents")

# Query and Rerank
query = "Which country has the second largest population"
results = rerank_pipeline.run({"retriever": {"query": query}, "reranker": {"query": query, "top_k": 3}})
print(results)

带元数据的完整示例

import os
from datasets import load_dataset
from haystack import Pipeline, Document
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.components.writers import DocumentWriter
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from mixedbread_ai_haystack import MixedbreadAIDocumentEmbedder, MixedbreadAITextEmbedder, MixedbreadAIReranker

# Set API Key
os.environ["MXBAI_API_KEY"] = "YOUR_API_KEY"

# Load the Dataset and Prepare Documents
ds = load_dataset("rajuptvs/ecommerce_products_clip")
documents = [
    Document(
        id=str(i),
        content=data["Description"], meta={
        "name": data["Product_name"],
        "price": data["Price"],
        "colors": data["colors"],
        "pattern": data["Pattern"],
        "extra": data["Other Details"]
    }) for i, data in enumerate(ds["train"])
]
meta_fields = documents[0].meta.keys()

# Define the Components
document_store = InMemoryDocumentStore(embedding_similarity_function="cosine")
document_writer = DocumentWriter(document_store=document_store)
embedding_retriever = InMemoryEmbeddingRetriever(document_store=document_store, top_k=20)

embed_model = "mixedbread-ai/mxbai-embed-large-v1"
reranking_model = "mixedbread-ai/mxbai-rerank-large-v1" 

text_embedder = MixedbreadAITextEmbedder(model=embed_model)
document_embedder = MixedbreadAIDocumentEmbedder(model=embed_model, max_concurrency=3, meta_fields_to_embed=meta_fields, show_progress_bar=True)
reranker = MixedbreadAIReranker(model=reranking_model, meta_fields_to_rank=meta_fields, top_k=5)

# Indexing Pipeline
indexing_pipeline = Pipeline()
indexing_pipeline.add_component(instance=document_embedder, name="document_embedder")
indexing_pipeline.add_component(instance=document_writer, name="document_writer")
indexing_pipeline.connect("document_embedder", "document_writer")

# Query Pipeline
query_pipeline = Pipeline()
query_pipeline.add_component(instance=text_embedder, name="text_embedder")
query_pipeline.add_component(instance=embedding_retriever, name="embedding_retriever")
query_pipeline.add_component(instance=reranker, name="reranker")
query_pipeline.connect("text_embedder", "embedding_retriever")
query_pipeline.connect("embedding_retriever.documents", "reranker.documents")

# Index the dataset
indexing_pipeline.run({"document_embedder": {"documents": documents}})

# Query to get results
query = "I am looking for a regular fit t-shirt in blue color. Ideally without any prints. What are my options?"
results = query_pipeline.run(
    {
        "text_embedder": {"text": query},
        "reranker": {"query": query}
    }
)
print(results["reranker"]["documents"])

集成：mixedbread ai

目录

概述

安装

使用

管道中的 Embedders

管道中的 Reranker

带元数据的完整示例