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

2026-06-30 16:40:46 +03:00
parent b186e2f395
commit 04e3b78a9c
2 changed files with 304 additions and 156 deletions
@@ -2,97 +2,94 @@

 ## Описание

-Это простая реализация поискового агента, использующего трансформерный энкодер для векторизации документов и запросов. Поиск осуществляется по косинусному сходству между векторами.
+Данный репозиторий содержит простую реализацию поискового агента, построенного с нуля с использованием глубоких нейронных сетей. Агент:

-## Структура проекта
+- **Преобразует** запросы и документы в TF‑IDF векторы.
+- **Обучается** на паре запрос‑документ с метками релевантности, используя логистическую регрессию (один линейный слой).
+- **Оценивает** точность и полноту на тестовом наборе.
+- **Возвращает** топ‑k наиболее релевантных документов для любого запроса.

-```
-src/
-├── index.py          # Основной код агента и API
-data/
-└── documents.txt     # Текстовый файл с документами (один документ на строку)
-README.md
-```
+### Технологии

-> **Важно**: файл `data/documents.txt` должен существовать и содержать хотя бы несколько строк текста. Если его нет, агент не запустится.
+- Python 3.8+
+- PyTorch
+- scikit‑learn
+- NumPy

 ## Установка

 ```bash
+# Клонируйте репозиторий
+git clone https://git.brojs.ru/kuzakhmetovartur/8.-samopisnyy-poiskovyy-agent-na-osnove.git
+cd 8.-samopisnyy-poiskovyy-agent-na-osnove
+
 # Создайте виртуальное окружение (рекомендуется)
 python -m venv venv
 source venv/bin/activate   # Windows: venv\Scripts\activate

 # Установите зависимости
-pip install -U pip
-pip install torch transformers fastapi uvicorn pydantic numpy
+pip install -r requirements.txt
 ```

-> Если у вас есть GPU, убедитесь, что установлена версия `torch` с поддержкой CUDA.
+`requirements.txt` содержит:

-## Запуск
+```
+torch>=1.7.0
+scikit-learn>=0.24
+numpy>=1.19
+```
+
+## Использование
+
+### Демонстрация

 ```bash
 python src/index.py
 ```

-Сервер будет доступен по адресу `http://0.0.0.0:8000`.
+Вы увидите вывод обучения, оценку точности/полноты и пример поиска.

-## API
-
-### POST `/search`
-
-**Запрос**
-
-```json
-{
-  "query": "пример запроса",
-  "top_k": 5
-}
-```
-
- `query` – строка запроса.
- `top_k` – количество возвращаемых документов (по умолчанию 5).
-
-**Ответ**
-
-```json
-{
-  "results": [
-    {
-      "document": "текст найденного документа",
-      "score": 0.8723
-    },
-    ...
-  ]
-}
-```
-
-## Пример использования
-
-```bash
-curl -X POST "http://0.0.0.0:8000/search" \
-     -H "Content-Type: application/json" \
-     -d '{"query":"машинное обучение", "top_k":3}'
-```
-
-## Как добавить документы
-
-1. Откройте файл `data/documents.txt`.
-2. Добавьте новые строки – каждая строка будет рассматриваться как отдельный документ.
-3. Перезапустите сервер, чтобы обновления вступили в силу.
-
-## Тесты
-
-Тестов в проекте нет, но вы можете быстро проверить работу:
+### Интеграция в свой проект

 ```python
 from src.index import SearchAgent

-agent = SearchAgent()
-print(agent.search("пример", top_k=3))
+documents = [
+    "Документ 1",
+    "Документ 2",
+    # ...
+]
+
+agent = SearchAgent(documents)
+
+# Обучаем
+queries = ["Какой вопрос?", "Что такое AI?"]
+relevance = [[0], [1]]  # индексы релевантных документов
+agent.train_agent(queries, relevance, epochs=10)
+
+# Оцениваем
+precision, recall = agent.evaluate(queries, relevance)
+
+# Поиск
+results = agent.search("Новый запрос", top_k=3)
+for doc, score in results:
+    print(f"{score:.4f} - {doc}")
 ```

+## Тесты
+
+Тесты находятся в `tests/` (если добавлены). Запуск:
+
+```bash
+pytest tests/
+```
+
+## Ограничения
+
+- Модель использует только один линейный слой, поэтому не может захватывать сложные взаимосвязи.
+- Размерность TF‑IDF может быть большой; для больших наборов данных стоит использовать более эффективные методы векторизации.
+- В примере используется синтетический набор данных; для реальных задач потребуется более крупный датасет и более сложная модель.
+
 ## Лицензия

 MIT License
@@ -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):
+
+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
-    top_k: int = 5
+            Query text.
+        top_k : int, default 3
+            Number of documents to return.

-class SearchResult(BaseModel):
-    document: str
-    score: float
-
-class SearchResponse(BaseModel):
-    results: List[SearchResult]
-
-class SearchAgent:
+        Returns
+        -------
+        List[Tuple[str, float]]
+            List of (document_text, score) tuples sorted by descending score.
        """
-    A simple deep‑agent search engine that uses a transformer encoder
-    to embed documents and queries, then ranks documents by cosine
-    similarity.
-    """
-    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 _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)
+        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)
-                batch_embeddings = sum_embeddings / sum_mask
-                embeddings.append(batch_embeddings.cpu().numpy())
-        return np.vstack(embeddings)
+            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 _embed_query(self, query: str) -> np.ndarray:
-        inputs = self.tokenizer([query],
-                                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)
-            query_embedding = sum_embeddings / sum_mask
-            return query_embedding.cpu().numpy()

-    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
-
-app = FastAPI(title="Deep Agent Search API")
-
-# Instantiate the agent once at startup
-agent = SearchAgent()
-
-@app.post("/search", response_model=SearchResponse)
-async def search_endpoint(request: SearchRequest):
+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()