imranjamal/Haystack

1# Import necessary libraries2import gradio as gr3from haystack.document_stores import InMemoryDocumentStore4from haystack.nodes import DensePassageRetriever, FARMReader5from haystack.pipelines import ExtractiveQAPipeline6 7# 1. Initialize Document Store8document_store = InMemoryDocumentStore(embedding_dim=768)9 10# 2. Add Documents11documents = [12    {"content": "Haystack is an open-source NLP framework for search.", "meta": {"source": "Introduction"}},13    {"content": "You can use Hugging Face models in Haystack pipelines.", "meta": {"source": "Hugging Face"}},14    {"content": "The DensePassageRetriever is a key component of Haystack.", "meta": {"source": "Retrievers"}}15]16document_store.write_documents(documents)17 18# 3. Set up Retriever19retriever = DensePassageRetriever(20    document_store=document_store,21    query_embedding_model="facebook/dpr-question_encoder-single-nq-base",22    passage_embedding_model="facebook/dpr-ctx_encoder-single-nq-base"23)24document_store.update_embeddings(retriever)25 26# 4. Set up Reader27reader = FARMReader(model_name_or_path="deepset/roberta-base-squad2", use_gpu=False)28 29# 5. Create QA Pipeline30qa_pipeline = ExtractiveQAPipeline(reader=reader, retriever=retriever)31 32# 6. Define Prediction Function33def ask_question(query):34    results = qa_pipeline.run(query=query, params={"Retriever": {"top_k": 3}, "Reader": {"top_k": 1}})35    if results["answers"]:36        answer = results["answers"][0].answer37        context = results["answers"][0].context38        source = results["answers"][0].meta.get("source", "Unknown Source")39        return f"**Answer:** {answer}\n\n**Context:** {context}\n\n**Source:** {source}"40    else:41        return "No relevant answer found. Please refine your question."42 43# 7. Set up Gradio Interface44interface = gr.Interface(45    fn=ask_question,46    inputs=gr.Textbox(lines=2, label="Ask a Question"),47    outputs="text",48    title="AI Search with Haystack",49    description="Ask any question about the content in the document set."50)51 52# 8. Launch Application53if __name__ == "__main__":54    interface.launch()55