diff --git a/README.md b/README.md index ce31ec3..1810d37 100644 --- a/README.md +++ b/README.md @@ -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 \ No newline at end of file diff --git a/src/index.py b/src/index.py index 6553620..d6e6cee 100644 --- a/src/index.py +++ b/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) \ No newline at end of file + demo() \ No newline at end of file