CoolFace
Apppublic

AlexGrig23/Unit_3_Agentic_RAG

sourceHugging Faceapache-2.0updated 2y agoView on Hugging Face
0likes
retriever.py54 linesDownload Raw Back to root
1from smolagents import Tool2from langchain_community.retrievers import BM25Retriever3from langchain.docstore.document import Document4import datasets5 6 7class GuestInfoRetrieverTool(Tool):8    name = "guest_info_retriever"9    description = "Retrieves detailed information about gala guests based on their name or relation."10    inputs = {11        "query": {12            "type": "string",13            "description": "The name or relation of the guest you want information about."14        }15    }16    output_type = "string"17 18    def __init__(self, docs):19        self.is_initialized = False20        self.retriever = BM25Retriever.from_documents(docs)21       22 23    def forward(self, query: str):24        results = self.retriever.get_relevant_documents(query)25        if results:26            return "\n\n".join([doc.page_content for doc in results[:3]])27        else:28            return "No matching guest information found."29 30 31def load_guest_dataset():32    # Load the dataset33    guest_dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")34 35    # Convert dataset entries into Document objects36    docs = [37        Document(38            page_content="\n".join([39                f"Name: {guest['name']}",40                f"Relation: {guest['relation']}",41                f"Description: {guest['description']}",42                f"Email: {guest['email']}"43            ]),44            metadata={"name": guest["name"]}45        )46        for guest in guest_dataset47    ]48 49    # Return the tool50    return GuestInfoRetrieverTool(docs)51 52 53 54