feat: solution for '8. Самописный поисковый агент на основе deep agents from scratch'

src/index.py

@@ -1,110 +1,261 @@
-import os
-import json
+#!/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
-from transformers import AutoTokenizer, AutoModel
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-from typing import List
+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 SearchRequest(BaseModel):
-    query: str
-    top_k: int = 5
 
-class SearchResult(BaseModel):
-    document: str
-    score: float
-
-class SearchResponse(BaseModel):
-    results: List[SearchResult]
-
-class SearchAgent:
+class SearchAgent(nn.Module):
     """
-    A simple deep‑agent search engine that uses a transformer encoder
-    to embed documents and queries, then ranks documents by cosine
-    similarity.
+    A simple search agent based on a linear classifier.
     """
-    def __init__(self,
-                 data_path: str = "data/documents.txt",
-                 model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModel.from_pretrained(model_name).to(self.device)
-        self.documents = self._load_documents(data_path)
-        self.embeddings = self._embed_documents(self.documents)
 
-    def _load_documents(self, path: str) -> List[str]:
-        if not os.path.exists(path):
-            raise FileNotFoundError(f"Data file not found: {path}")
-        with open(path, "r", encoding="utf-8") as f:
-            docs = [line.strip() for line in f if line.strip()]
-        return docs
+    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")
 
-    def _embed_documents(self, docs: List[str]) -> np.ndarray:
-        batch_size = 32
-        embeddings = []
-        for i in range(0, len(docs), batch_size):
-            batch = docs[i:i+batch_size]
-            inputs = self.tokenizer(batch,
-                                    padding=True,
-                                    truncation=True,
-                                    return_tensors="pt").to(self.device)
-            with torch.no_grad():
-                outputs = self.model(**inputs)
-                token_embeddings = outputs.last_hidden_state
-                attention_mask = inputs["attention_mask"].unsqueeze(-1)
-                sum_embeddings = torch.sum(token_embeddings * attention_mask, dim=1)
-                sum_mask = torch.clamp(attention_mask.sum(dim=1), min=1e-9)
-                batch_embeddings = sum_embeddings / sum_mask
-                embeddings.append(batch_embeddings.cpu().numpy())
-        return np.vstack(embeddings)
+        # 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)
 
-    def _embed_query(self, query: str) -> np.ndarray:
-        inputs = self.tokenizer([query],
-                                padding=True,
-                                truncation=True,
-                                return_tensors="pt").to(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():
-            outputs = self.model(**inputs)
-            token_embeddings = outputs.last_hidden_state
-            attention_mask = inputs["attention_mask"].unsqueeze(-1)
-            sum_embeddings = torch.sum(token_embeddings * attention_mask, dim=1)
-            sum_mask = torch.clamp(attention_mask.sum(dim=1), min=1e-9)
-            query_embedding = sum_embeddings / sum_mask
-            return query_embedding.cpu().numpy()
+            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()
 
-    def search(self, query: str, top_k: int = 5) -> List[SearchResult]:
-        query_emb = self._embed_query(query)
-        dot = np.dot(self.embeddings, query_emb.T).squeeze()
-        norms = np.linalg.norm(self.embeddings, axis=1) * np.linalg.norm(query_emb)
-        similarities = dot / norms
-        top_indices = np.argsort(similarities)[::-1][:top_k]
-        results = [SearchResult(document=self.documents[idx],
-                                score=float(similarities[idx]))
-                   for idx in top_indices]
-        return results
+        # Compute precision and recall
+        total_relevant = 0
+        total_predicted = 0
+        total_correct = 0
 
-app = FastAPI(title="Deep Agent Search API")
+        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))
 
-# Instantiate the agent once at startup
-agent = SearchAgent()
+        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
 
-@app.post("/search", response_model=SearchResponse)
-async def search_endpoint(request: SearchRequest):
+    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():
     """
-    Search endpoint that accepts a JSON payload:
-    {
-        "query": "your search query",
-        "top_k": 5
-    }
-    Returns the top_k most relevant documents.
+    Demo usage of the SearchAgent with a tiny synthetic dataset.
     """
-    try:
-        results = agent.search(request.query, request.top_k)
-        return SearchResponse(results=results)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+    # 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__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    demo()