261 lines
9.0 KiB
Python
261 lines
9.0 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Simple Deep Search Agent
|
||
========================
|
||
|
||
This module implements a minimal deep learning based search agent.
|
||
It uses a TF‑IDF vectorizer to transform documents and queries into
|
||
feature vectors and a single linear layer (logistic regression)
|
||
to predict relevance scores. The agent can be trained on a small
|
||
synthetic dataset and used to retrieve the top‑k most relevant
|
||
documents for a given query.
|
||
|
||
Author: Artur Kuzakhmetov
|
||
Date: 30.06.2026
|
||
"""
|
||
|
||
import sys
|
||
from typing import List, Tuple
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.optim as optim
|
||
from sklearn.feature_extraction.text import TfidfVectorizer
|
||
from sklearn.model_selection import train_test_split
|
||
|
||
|
||
class SearchAgent(nn.Module):
|
||
"""
|
||
A simple search agent based on a linear classifier.
|
||
"""
|
||
|
||
def __init__(self, documents: List[str], device: torch.device = None):
|
||
"""
|
||
Parameters
|
||
----------
|
||
documents : List[str]
|
||
List of document texts.
|
||
device : torch.device, optional
|
||
Device to run the model on. Defaults to CUDA if available.
|
||
"""
|
||
super().__init__()
|
||
self.documents = documents
|
||
self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
||
# Fit TF‑IDF vectorizer on documents
|
||
self.vectorizer = TfidfVectorizer()
|
||
self.doc_vectors = self.vectorizer.fit_transform(self.documents).toarray()
|
||
self.doc_vectors = torch.tensor(self.doc_vectors, dtype=torch.float32, device=self.device)
|
||
|
||
# Linear layer: input_dim -> 1 (relevance score)
|
||
self.linear = nn.Linear(self.doc_vectors.shape[1], 1).to(self.device)
|
||
|
||
def forward(self, query_vec: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Forward pass: compute relevance scores for all documents given a query vector.
|
||
|
||
Parameters
|
||
----------
|
||
query_vec : torch.Tensor
|
||
Tensor of shape (1, feature_dim).
|
||
|
||
Returns
|
||
-------
|
||
torch.Tensor
|
||
Tensor of shape (num_documents,) with relevance scores.
|
||
"""
|
||
# Compute dot product between query and each document vector
|
||
scores = torch.matmul(self.doc_vectors, query_vec.t()).squeeze(1)
|
||
return scores
|
||
|
||
def train_agent(
|
||
self,
|
||
queries: List[str],
|
||
labels: List[List[int]],
|
||
epochs: int = 10,
|
||
lr: float = 0.01,
|
||
batch_size: int = 4,
|
||
verbose: bool = True,
|
||
) -> None:
|
||
"""
|
||
Train the agent on query‑document relevance pairs.
|
||
|
||
Parameters
|
||
----------
|
||
queries : List[str]
|
||
List of query texts.
|
||
labels : List[List[int]]
|
||
List of relevance labels for each query. Each inner list contains
|
||
indices of relevant documents (0‑based).
|
||
epochs : int, default 10
|
||
Number of training epochs.
|
||
lr : float, default 0.01
|
||
Learning rate.
|
||
batch_size : int, default 4
|
||
Batch size.
|
||
verbose : bool, default True
|
||
Whether to print training progress.
|
||
"""
|
||
# Vectorize queries
|
||
query_vectors = self.vectorizer.transform(queries).toarray()
|
||
query_vectors = torch.tensor(query_vectors, dtype=torch.float32, device=self.device)
|
||
|
||
# Prepare training data
|
||
# For each query, create a target vector of relevance scores (1 for relevant, 0 otherwise)
|
||
targets = []
|
||
for rel_indices in labels:
|
||
target = torch.zeros(self.doc_vectors.shape[0], device=self.device)
|
||
target[rel_indices] = 1.0
|
||
targets.append(target)
|
||
targets = torch.stack(targets) # shape: (num_queries, num_documents)
|
||
|
||
# Loss and optimizer
|
||
criterion = nn.BCEWithLogitsLoss()
|
||
optimizer = optim.Adam(self.parameters(), lr=lr)
|
||
|
||
dataset = torch.utils.data.TensorDataset(query_vectors, targets)
|
||
loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
|
||
|
||
self.train()
|
||
for epoch in range(1, epochs + 1):
|
||
epoch_loss = 0.0
|
||
for batch_q, batch_t in loader:
|
||
optimizer.zero_grad()
|
||
outputs = self.forward(batch_q) # shape: (batch_size, num_documents)
|
||
loss = criterion(outputs, batch_t)
|
||
loss.backward()
|
||
optimizer.step()
|
||
epoch_loss += loss.item() * batch_q.size(0)
|
||
epoch_loss /= len(dataset)
|
||
if verbose:
|
||
print(f"Epoch {epoch}/{epochs} - Loss: {epoch_loss:.4f}")
|
||
|
||
def evaluate(
|
||
self,
|
||
queries: List[str],
|
||
labels: List[List[int]],
|
||
threshold: float = 0.5,
|
||
) -> Tuple[float, float]:
|
||
"""
|
||
Evaluate the agent on a test set.
|
||
|
||
Parameters
|
||
----------
|
||
queries : List[str]
|
||
List of query texts.
|
||
labels : List[List[int]]
|
||
List of relevance labels for each query.
|
||
threshold : float, default 0.5
|
||
Threshold to convert scores to binary predictions.
|
||
|
||
Returns
|
||
-------
|
||
Tuple[float, float]
|
||
(precision, recall)
|
||
"""
|
||
self.eval()
|
||
with torch.no_grad():
|
||
query_vectors = self.vectorizer.transform(queries).toarray()
|
||
query_vectors = torch.tensor(query_vectors, dtype=torch.float32, device=self.device)
|
||
outputs = self.forward(query_vectors) # shape: (num_queries, num_documents)
|
||
preds = (outputs > threshold).int()
|
||
|
||
# Compute precision and recall
|
||
total_relevant = 0
|
||
total_predicted = 0
|
||
total_correct = 0
|
||
|
||
for i, rel_indices in enumerate(labels):
|
||
pred_indices = preds[i].nonzero(as_tuple=True)[0].cpu().numpy().tolist()
|
||
total_relevant += len(rel_indices)
|
||
total_predicted += len(pred_indices)
|
||
total_correct += len(set(pred_indices) & set(rel_indices))
|
||
|
||
precision = total_correct / total_predicted if total_predicted > 0 else 0.0
|
||
recall = total_correct / total_relevant if total_relevant > 0 else 0.0
|
||
return precision, recall
|
||
|
||
def search(self, query: str, top_k: int = 3) -> List[Tuple[str, float]]:
|
||
"""
|
||
Retrieve top‑k documents for a given query.
|
||
|
||
Parameters
|
||
----------
|
||
query : str
|
||
Query text.
|
||
top_k : int, default 3
|
||
Number of documents to return.
|
||
|
||
Returns
|
||
-------
|
||
List[Tuple[str, float]]
|
||
List of (document_text, score) tuples sorted by descending score.
|
||
"""
|
||
self.eval()
|
||
with torch.no_grad():
|
||
q_vec = self.vectorizer.transform([query]).toarray()
|
||
q_vec = torch.tensor(q_vec, dtype=torch.float32, device=self.device)
|
||
scores = self.forward(q_vec).cpu().numpy().flatten()
|
||
top_indices = np.argsort(scores)[::-1][:top_k]
|
||
return [(self.documents[i], float(scores[i])) for i in top_indices]
|
||
|
||
|
||
def demo():
|
||
"""
|
||
Demo usage of the SearchAgent with a tiny synthetic dataset.
|
||
"""
|
||
# Sample documents
|
||
docs = [
|
||
"Deep learning models can learn complex patterns from data.",
|
||
"Search engines index web pages to provide relevant results.",
|
||
"Natural language processing enables computers to understand text.",
|
||
"Python is a popular programming language for data science.",
|
||
"Machine learning algorithms improve over time with more data.",
|
||
"Artificial intelligence encompasses machine learning and deep learning.",
|
||
"Information retrieval is a key component of search systems.",
|
||
"Neural networks consist of layers of interconnected nodes.",
|
||
"Data preprocessing is essential before training models.",
|
||
"Evaluation metrics help assess model performance.",
|
||
]
|
||
|
||
# Sample queries and relevance labels (indices of relevant docs)
|
||
queries = [
|
||
"What is deep learning?",
|
||
"How do search engines work?",
|
||
"Explain natural language processing.",
|
||
"Why use Python for data science?",
|
||
"What is machine learning?",
|
||
]
|
||
|
||
relevance = [
|
||
[0, 5], # Relevant docs for query 0
|
||
[1, 6], # Relevant docs for query 1
|
||
[2, 7], # Relevant docs for query 2
|
||
[3, 9], # Relevant docs for query 3
|
||
[4, 5], # Relevant docs for query 4
|
||
]
|
||
|
||
# Split into train/test
|
||
train_q, test_q, train_rel, test_rel = train_test_split(
|
||
queries, relevance, test_size=0.4, random_state=42
|
||
)
|
||
|
||
agent = SearchAgent(docs)
|
||
print("Training agent...")
|
||
agent.train_agent(train_q, train_rel, epochs=20, lr=0.01, verbose=True)
|
||
|
||
print("\nEvaluating agent...")
|
||
precision, recall = agent.evaluate(test_q, test_rel)
|
||
print(f"Precision: {precision:.2f}, Recall: {recall:.2f}")
|
||
|
||
# Search example
|
||
query = "Tell me about deep learning and AI."
|
||
print(f"\nSearching for: '{query}'")
|
||
results = agent.search(query, top_k=5)
|
||
for doc, score in results:
|
||
print(f"Score: {score:.4f} | Doc: {doc}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
demo() |