Spaces:
Runtime error
Runtime error
File size: 2,013 Bytes
89efac4 7595a2c 89efac4 7595a2c 89efac4 7595a2c 89efac4 7595a2c 89efac4 62d1a65 89efac4 7595a2c 89efac4 7595a2c 89efac4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
"""Guest Information Retrieval Toll"""
from smolagents import Tool
from langchain_community.retrievers import BM25Retriever
from langchain.docstore.document import Document
import datasets
class GuestInfoRetrieverTool(Tool):
"""Derived Class for Guest Information Retrieval Tool"""
name = "guest_info_retriever"
description = "Retrieves detailed information about gala guests based on their name or relation." # pylint: disable=line-too-long
inputs = {
"query": {
"type": "string",
"description": "The name or relation of the guest you want information about.",
}
}
output_type = "string"
def __init__(self, docs): # pylint: disable=super-init-not-called
self.is_initialized = False
self.retriever = BM25Retriever.from_documents(docs)
def forward(self, query: str): # pylint: disable=arguments-differ
results = self.retriever.get_relevant_documents(query)
if results:
return "\n\n".join([doc.page_content for doc in results[:3]])
return "No matching guest information found."
def guest_info_retriever_factory():
"""Get Guest Information Retrieval Tool"""
hf_datasets = [
"agents-course/unit3-invitees",
"Data-Gem/agents-course-unit3-invitees-expanded",
]
guest_dataset = []
for dataset_name in hf_datasets:
new_dataset = datasets.load_dataset(dataset_name, split="train")
guest_dataset.append(new_dataset)
# Convert dataset entries into Document objects
docs = [
Document(
page_content="\n".join(
[
f"Name: {guest['name']}",
f"Relation: {guest['relation']}",
f"Description: {guest['description']}",
f"Email: {guest['email']}",
]
),
metadata={"name": guest["name"]},
)
for guest in guest_dataset
]
return GuestInfoRetrieverTool(docs)
|