haystack/tutorials/Tutorial14_Query_Classifier.py

import os
from subprocess import Popen, PIPE, STDOUT
from haystack.utils import fetch_archive_from_http, convert_files_to_dicts, clean_wiki_text, launch_es, print_answers
from haystack.pipelines import Pipeline, RootNode
from haystack.document_stores import ElasticsearchDocumentStore
from haystack.nodes import (
    ElasticsearchRetriever,
    DensePassageRetriever,
    FARMReader,
    TransformersQueryClassifier,
    SklearnQueryClassifier,
)


def tutorial14_query_classifier():

    # Download and prepare data - 517 Wikipedia articles for Game of Thrones
    doc_dir = "data/article_txt_got"
    s3_url = "https://s3.eu-central-1.amazonaws.com/deepset.ai-farm-qa/datasets/documents/wiki_gameofthrones_txt.zip"
    fetch_archive_from_http(url=s3_url, output_dir=doc_dir)

    # convert files to dicts containing documents that can be indexed to our datastore
    got_dicts = convert_files_to_dicts(dir_path=doc_dir, clean_func=clean_wiki_text, split_paragraphs=True)

    # Initialize DocumentStore and index documents
    launch_es()
    document_store = ElasticsearchDocumentStore()
    document_store.delete_documents()
    document_store.write_documents(got_dicts)

    # Initialize Sparse retriever
    es_retriever = ElasticsearchRetriever(document_store=document_store)

    # Initialize dense retriever
    dpr_retriever = DensePassageRetriever(document_store)
    document_store.update_embeddings(dpr_retriever, update_existing_embeddings=False)

    reader = FARMReader(model_name_or_path="deepset/roberta-base-squad2")

    print()
    print("Sklearn keyword classifier")
    print("==========================")
    # Here we build the pipeline
    sklearn_keyword_classifier = Pipeline()
    sklearn_keyword_classifier.add_node(component=SklearnQueryClassifier(), name="QueryClassifier", inputs=["Query"])
    sklearn_keyword_classifier.add_node(
        component=dpr_retriever, name="DPRRetriever", inputs=["QueryClassifier.output_1"]
    )
    sklearn_keyword_classifier.add_node(component=es_retriever, name="ESRetriever", inputs=["QueryClassifier.output_2"])
    sklearn_keyword_classifier.add_node(component=reader, name="QAReader", inputs=["ESRetriever", "DPRRetriever"])
    sklearn_keyword_classifier.draw("pipeline_classifier.png")

    # Run only the dense retriever on the full sentence query
    res_1 = sklearn_keyword_classifier.run(
        query="Who is the father of Arya Stark?",
    )
    print("\n===============================")
    print("DPR Results" + "\n" + "=" * 15)
    print_answers(res_1, details="minimum")

    # Run only the sparse retriever on a keyword based query
    res_2 = sklearn_keyword_classifier.run(
        query="arya stark father",
    )
    print("\n===============================")
    print("ES Results" + "\n" + "=" * 15)
    print_answers(res_2, details="minimum")

    # Run only the dense retriever on the full sentence query
    res_3 = sklearn_keyword_classifier.run(
        query="which country was jon snow filmed ?",
    )
    print("\n===============================")
    print("DPR Results" + "\n" + "=" * 15)
    print_answers(res_3, details="minimum")

    # Run only the sparse retriever on a keyword based query
    res_4 = sklearn_keyword_classifier.run(
        query="jon snow country",
    )
    print("\n===============================")
    print("ES Results" + "\n" + "=" * 15)
    print_answers(res_4, details="minimum")

    # Run only the dense retriever on the full sentence query
    res_5 = sklearn_keyword_classifier.run(
        query="who are the younger brothers of arya stark ?",
    )
    print("\n===============================")
    print("DPR Results" + "\n" + "=" * 15)
    print_answers(res_5, details="minimum")

    # Run only the sparse retriever on a keyword based query
    res_6 = sklearn_keyword_classifier.run(
        query="arya stark younger brothers",
    )
    print("\n===============================")
    print("ES Results" + "\n" + "=" * 15)
    print_answers(res_6, details="minimum")

    print()
    print("Transformer keyword classifier")
    print("==============================")
    # Here we build the pipeline
    transformer_keyword_classifier = Pipeline()
    transformer_keyword_classifier.add_node(
        component=TransformersQueryClassifier(), name="QueryClassifier", inputs=["Query"]
    )
    transformer_keyword_classifier.add_node(
        component=dpr_retriever, name="DPRRetriever", inputs=["QueryClassifier.output_1"]
    )
    transformer_keyword_classifier.add_node(
        component=es_retriever, name="ESRetriever", inputs=["QueryClassifier.output_2"]
    )
    transformer_keyword_classifier.add_node(component=reader, name="QAReader", inputs=["ESRetriever", "DPRRetriever"])
    transformer_keyword_classifier.draw("pipeline_classifier.png")

    # Run only the dense retriever on the full sentence query
    res_1 = transformer_keyword_classifier.run(
        query="Who is the father of Arya Stark?",
    )
    print("\n===============================")
    print("DPR Results" + "\n" + "=" * 15)
    print_answers(res_1, details="minimum")

    # Run only the sparse retriever on a keyword based query
    res_2 = transformer_keyword_classifier.run(
        query="arya stark father",
    )
    print("\n===============================")
    print("ES Results" + "\n" + "=" * 15)
    print_answers(res_2, details="minimum")

    # Run only the dense retriever on the full sentence query
    res_3 = transformer_keyword_classifier.run(
        query="which country was jon snow filmed ?",
    )
    print("\n===============================")
    print("DPR Results" + "\n" + "=" * 15)
    print_answers(res_3, details="minimum")

    # Run only the sparse retriever on a keyword based query
    res_4 = transformer_keyword_classifier.run(
        query="jon snow country",
    )
    print("\n===============================")
    print("ES Results" + "\n" + "=" * 15)
    print_answers(res_4, details="minimum")

    # Run only the dense retriever on the full sentence query
    res_5 = transformer_keyword_classifier.run(
        query="who are the younger brothers of arya stark ?",
    )
    print("\n===============================")
    print("DPR Results" + "\n" + "=" * 15)
    print_answers(res_5, details="minimum")

    # Run only the sparse retriever on a keyword based query
    res_6 = transformer_keyword_classifier.run(
        query="arya stark younger brothers",
    )
    print("\n===============================")
    print("ES Results" + "\n" + "=" * 15)
    print_answers(res_6, details="minimum")

    print()
    print("Transformer question classifier")
    print("===============================")

    # Here we build the pipeline
    transformer_question_classifier = Pipeline()
    transformer_question_classifier.add_node(component=dpr_retriever, name="DPRRetriever", inputs=["Query"])
    transformer_question_classifier.add_node(
        component=TransformersQueryClassifier(model_name_or_path="shahrukhx01/question-vs-statement-classifier"),
        name="QueryClassifier",
        inputs=["DPRRetriever"],
    )
    transformer_question_classifier.add_node(component=reader, name="QAReader", inputs=["QueryClassifier.output_1"])
    transformer_question_classifier.draw("question_classifier.png")

    # Run only the QA reader on the question query
    res_1 = transformer_question_classifier.run(
        query="Who is the father of Arya Stark?",
    )
    print("\n===============================")
    print("DPR Results" + "\n" + "=" * 15)
    print_answers(res_1, details="minimum")

    # Show only DPR results
    res_2 = transformer_question_classifier.run(
        query="Arya Stark was the daughter of a Lord.",
    )
    print("\n===============================")
    print("ES Results" + "\n" + "=" * 15)
    print_answers(res_2, details="minimum")

    # Here we create the keyword vs question/statement query classifier

    queries = [
        "arya stark father",
        "jon snow country",
        "who is the father of arya stark",
        "which country was jon snow filmed?",
    ]

    keyword_classifier = TransformersQueryClassifier()

    for query in queries:
        result = keyword_classifier.run(query=query)
        if result[1] == "output_1":
            category = "question/statement"
        else:
            category = "keyword"

        print(f"Query: {query}, raw_output: {result}, class: {category}")

    # Here we create the question vs statement query classifier

    queries = [
        "Lord Eddard was the father of Arya Stark.",
        "Jon Snow was filmed in United Kingdom.",
        "who is the father of arya stark?",
        "Which country was jon snow filmed in?",
    ]

    question_classifier = TransformersQueryClassifier(model_name_or_path="shahrukhx01/question-vs-statement-classifier")

    for query in queries:
        result = question_classifier.run(query=query)
        if result[1] == "output_1":
            category = "question"
        else:
            category = "statement"

        print(f"Query: {query}, raw_output: {result}, class: {category}")


if __name__ == "__main__":
    tutorial14_query_classifier()

# This Haystack script was made with love by deepset in Berlin, Germany
# Haystack: https://github.com/deepset-ai/haystack
# deepset: https://deepset.ai/