116 lines
3.3 KiB
Python
116 lines
3.3 KiB
Python
"""
|
|
Main entry point for the deep agent search application.
|
|
|
|
Demonstrates:
|
|
- Creating a virtual file system.
|
|
- Adding a virtual file with sample data.
|
|
- Performing a search query using a simple deep agent.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from typing import List
|
|
|
|
import numpy as np
|
|
import torch
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
|
|
from virtual_file_system import VirtualFileSystem, VirtualFile
|
|
|
|
|
|
class SearchAgent:
|
|
"""
|
|
Simple search agent that ranks lines from a virtual file based on
|
|
cosine similarity between TF-IDF vectors of the query and the lines.
|
|
"""
|
|
|
|
def __init__(self, vfs: VirtualFileSystem):
|
|
self.vfs = vfs
|
|
|
|
def search(self, file_name: str, query: str, top_k: int = 5) -> List[str]:
|
|
"""
|
|
Search for the most relevant lines in the specified virtual file.
|
|
|
|
Parameters
|
|
----------
|
|
file_name : str
|
|
Name of the virtual file to search.
|
|
query : str
|
|
Search query string.
|
|
top_k : int, optional
|
|
Number of top results to return.
|
|
|
|
Returns
|
|
-------
|
|
List[str]
|
|
List of the most relevant lines.
|
|
"""
|
|
vf = self.vfs.get_file(file_name)
|
|
data = vf.read().decode("utf-8")
|
|
lines = [line.strip() for line in data.splitlines() if line.strip()]
|
|
if not lines:
|
|
return []
|
|
|
|
# Vectorize lines and query
|
|
vectorizer = TfidfVectorizer()
|
|
doc_vectors = vectorizer.fit_transform(lines).toarray()
|
|
query_vec = vectorizer.transform([query]).toarray()
|
|
|
|
# Convert to torch tensors for similarity calculation
|
|
doc_tensors = torch.tensor(doc_vectors, dtype=torch.float32)
|
|
query_tensor = torch.tensor(query_vec, dtype=torch.float32)
|
|
|
|
# Normalize vectors
|
|
doc_norm = doc_tensors / doc_tensors.norm(dim=1, keepdim=True)
|
|
query_norm = query_tensor / query_tensor.norm()
|
|
|
|
# Cosine similarity
|
|
similarities = torch.matmul(doc_norm, query_norm.t()).squeeze()
|
|
|
|
# Get top_k indices
|
|
top_indices = similarities.topk(top_k).indices.tolist()
|
|
return [lines[i] for i in top_indices]
|
|
|
|
|
|
def main(argv: List[str]) -> None:
|
|
"""
|
|
Example usage of the virtual file system and search agent.
|
|
|
|
Creates a virtual file with sample text and performs a search query.
|
|
"""
|
|
vfs = VirtualFileSystem()
|
|
|
|
# Sample data: a small collection of sentences
|
|
sample_text = """\
|
|
Deep learning has revolutionized many fields.
|
|
Neural networks can approximate complex functions.
|
|
PyTorch provides dynamic computation graphs.
|
|
Scikit-learn offers a wide range of machine learning tools.
|
|
Numpy is essential for numerical operations.
|
|
"""
|
|
|
|
# Create a virtual file
|
|
vf = vfs.create_file("sample.txt", sample_text.encode("utf-8"))
|
|
|
|
# Instantiate the search agent
|
|
agent = SearchAgent(vfs)
|
|
|
|
# Perform a search query
|
|
query = "neural networks"
|
|
results = agent.search("sample.txt", query, top_k=3)
|
|
|
|
print(f"Search results for query: '{query}'")
|
|
for idx, line in enumerate(results, 1):
|
|
print(f"{idx}. {line}")
|
|
|
|
# Demonstrate unload
|
|
vf.unload()
|
|
try:
|
|
vf.read()
|
|
except RuntimeError as e:
|
|
print(f"After unload: {e}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main(sys.argv[1:]) |