jarguello76 commited on
Commit
1b478b2
·
verified ·
1 Parent(s): 8e3585f

Delete retriever.py

Browse files
Files changed (1) hide show
  1. retriever.py +0 -47
retriever.py DELETED
@@ -1,47 +0,0 @@
1
- from smolagents import Tool
2
- from langchain_community.retrievers import BM25Retriever
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."
9
- inputs = {
10
- "query": {
11
- "type": "string",
12
- "description": "The name or relation of the guest you want information about."
13
- }
14
- }
15
- output_type = "string"
16
-
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)
23
- if results:
24
- return "\n\n".join([doc.page_content for doc in results[:3]])
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")
31
-
32
- # Convert dataset entries into Document objects
33
- docs = [
34
- Document(
35
- page_content="\n".join([
36
- f"Name: {guest['name']}",
37
- f"Relation: {guest['relation']}",
38
- f"Description: {guest['description']}",
39
- f"Email: {guest['email']}"
40
- ]),
41
- metadata={"name": guest["name"]}
42
- )
43
- for guest in guest_dataset
44
- ]
45
-
46
- # Return the tool
47
- return GuestInfoRetrieverTool(docs)