Update retriever.py
Browse files- retriever.py +1 -4
retriever.py
CHANGED
@@ -3,7 +3,6 @@ from langchain_community.retrievers import BM25Retriever
|
|
3 |
from langchain.docstore.document import Document
|
4 |
import datasets
|
5 |
|
6 |
-
|
7 |
class GuestInfoRetrieverTool(Tool):
|
8 |
name = "guest_info_retriever"
|
9 |
description = "Retrieves detailed information about gala guests based on their name or relation."
|
@@ -18,7 +17,6 @@ class GuestInfoRetrieverTool(Tool):
|
|
18 |
def __init__(self, docs):
|
19 |
self.is_initialized = False
|
20 |
self.retriever = BM25Retriever.from_documents(docs)
|
21 |
-
|
22 |
|
23 |
def forward(self, query: str):
|
24 |
results = self.retriever.get_relevant_documents(query)
|
@@ -27,7 +25,6 @@ class GuestInfoRetrieverTool(Tool):
|
|
27 |
else:
|
28 |
return "No matching guest information found."
|
29 |
|
30 |
-
|
31 |
def load_guest_dataset():
|
32 |
# Load the dataset
|
33 |
guest_dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")
|
@@ -47,4 +44,4 @@ def load_guest_dataset():
|
|
47 |
]
|
48 |
|
49 |
# Return the tool
|
50 |
-
return GuestInfoRetrieverTool(docs)
|
|
|
3 |
from langchain.docstore.document import Document
|
4 |
import datasets
|
5 |
|
|
|
6 |
class GuestInfoRetrieverTool(Tool):
|
7 |
name = "guest_info_retriever"
|
8 |
description = "Retrieves detailed information about gala guests based on their name or relation."
|
|
|
17 |
def __init__(self, docs):
|
18 |
self.is_initialized = False
|
19 |
self.retriever = BM25Retriever.from_documents(docs)
|
|
|
20 |
|
21 |
def forward(self, query: str):
|
22 |
results = self.retriever.get_relevant_documents(query)
|
|
|
25 |
else:
|
26 |
return "No matching guest information found."
|
27 |
|
|
|
28 |
def load_guest_dataset():
|
29 |
# Load the dataset
|
30 |
guest_dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")
|
|
|
44 |
]
|
45 |
|
46 |
# Return the tool
|
47 |
+
return GuestInfoRetrieverTool(docs)
|