Integration: INSTRUCTOR Embedders

pip install instructor-embedders-haystack

from instructor_embedders.instructor_text_embedder import InstructorTextEmbedder

# Example text from the Amazon Reviews Polarity Dataset (https://huggingface.co/datasets/amazon_polarity)
text = "It clearly says online this will work on a Mac OS system. The disk comes and it does not, only Windows. Do Not order this if you have a Mac!!"
instruction = (
    "Represent the Amazon comment for classifying the sentence as positive or negative"
)

text_embedder = InstructorTextEmbedder(
    model_name_or_path="hkunlp/instructor-base", instruction=instruction,
    device="cpu"
)

from instructor_embedders.instructor_document_embedder import InstructorDocumentEmbedder
from haystack.preview.dataclasses import Document


doc_embedding_instruction = "Represent the Medical Document for retrieval:"

doc_embedder = InstructorDocumentEmbedder(
    model_name_or_path="hkunlp/instructor-base",
    instruction=doc_embedding_instruction,
    batch_size=32,
    device="cpu",
)

doc_embedder.warm_up()

# Text taken from PubMed QA Dataset (https://huggingface.co/datasets/pubmed_qa)
document_list = [
    Document(
        text="Oxidative stress generated within inflammatory joints can produce autoimmune phenomena and joint destruction. Radical species with oxidative activity, including reactive nitrogen species, represent mediators of inflammation and cartilage damage.",
        metadata={
            "pubid": "25,445,628",
            "long_answer": "yes",
        },
    ),
    Document(
        text="Plasma levels of pancreatic polypeptide (PP) rise upon food intake. Although other pancreatic islet hormones, such as insulin and glucagon, have been extensively investigated, PP secretion and actions are still poorly understood.",
        metadata={
            "pubid": "25,445,712",
            "long_answer": "yes",
        },
    ),
    Document(
        text="Disturbed sleep is associated with mood disorders. Both depression and insomnia may increase the risk of disability retirement. The longitudinal links among insomnia, depression and work incapacity are poorly known.",
        metadata={
            "pubid": "25,451,441",
            "long_answer": "yes",
        },
    ),
]

result = doc_embedder.run(document_list)
print(f"Document Text: {result['documents'][0].text}")
print(f"Document Embedding: {result['documents'][0].embedding}")
print(f"Embedding Dimension: {len(result['documents'][0].embedding)}")

# Import necessary modules and classes
from haystack.preview.document_stores import InMemoryDocumentStore
from haystack.preview.dataclasses import Document
from haystack.preview import Pipeline
from haystack.preview.components.writers import DocumentWriter
from haystack.preview.components.retrievers import InMemoryEmbeddingRetriever
from datasets import load_dataset

# Import custom INSTRUCTOR Embedders
from instructor_embedders.instructor_document_embedder import InstructorDocumentEmbedder
from instructor_embedders.instructor_text_embedder import InstructorTextEmbedder

# Initialize a InMemoryDocumentStore, which will be used to store and retrieve documents
# It uses cosine similarity for document embeddings comparison
doc_store = InMemoryDocumentStore(embedding_similarity_function="cosine")

# Define an instruction for document embedding
doc_embedding_instruction = "Represent the News Article for retrieval:"
# Create an InstructorDocumentEmbedder instance with specified parameters
doc_embedder = InstructorDocumentEmbedder(
    model_name_or_path="hkunlp/instructor-base",
    instruction=doc_embedding_instruction,
    batch_size=32,
    device="cpu",
)
# Warm up the embedder (loading the pre-trained model)
doc_embedder.warm_up()

# Create an indexing pipeline
indexing_pipeline = Pipeline()
# Add the document embedder component to the pipeline
indexing_pipeline.add_component(instance=doc_embedder, name="DocEmbedder")
# Add a DocumentWriter component to the pipeline that writes documents to the Document Store
indexing_pipeline.add_component(
    instance=DocumentWriter(document_store=doc_store), name="DocWriter"
)
# Connect the output of DocEmbedder to the input of DocWriter
indexing_pipeline.connect(connect_from="DocEmbedder", connect_to="DocWriter")

# Load the 'XSum' dataset from HuggingFace (https://huggingface.co/datasets/xsum)
dataset = load_dataset("xsum", split="train")

# Create Document objects from the dataset and add them to the document store using the indexing pipeline
docs = [
    Document(
        text=doc["document"],
        metadata={
            "summary": doc["summary"],
            "doc_id": doc["id"],
        },
    )
    for doc in dataset
]
indexing_pipeline.run({"DocEmbedder": {"documents": docs}})

# Print the first document and its embedding from the document store
print(doc_store.filter_documents()[0])
print(doc_store.filter_documents()[0].embedding)

# Define an instruction for query embedding
query_embedding_instruction = (
    "Represent the news question for retrieving supporting articles:"
)
# Create an InstructorTextEmbedder instance for query embedding
text_embedder = InstructorTextEmbedder(
    model_name_or_path="hkunlp/instructor-base",
    instruction=query_embedding_instruction,
    device="cpu",
)
# Load the text embedding model
text_embedder.warm_up()

# Create a query pipeline
query_pipeline = Pipeline()
# Add the text embedder component to the pipeline
query_pipeline.add_component("TextEmbedder", text_embedder)
# Add a InMemoryEmbeddingRetriever component to the pipeline that retrieves documents from the doc_store
query_pipeline.add_component(
    "Retriever", InMemoryEmbeddingRetriever(document_store=doc_store)
)
# Connect the output of TextEmbedder to the input of Retriever
query_pipeline.connect("TextEmbedder", "Retriever")

# Run the query pipeline with a sample query text
results = query_pipeline.run(
    {
        "TextEmbedder": {
            "text": "What were the concerns expressed by Jeanette Tate regarding the response to the flooding in Newton Stewart?"
        }
    }
)

# Print information about retrieved documents
for doc in results["Retriever"]["documents"]:
    print(f"Text:\n{doc.text[:150]}...\n")
    print(f"Metadata: {doc.metadata}")
    print(f"Score: {doc.score}")
    print("-" * 10 + "\n")