Minimal RAG example. Local LLM (Qwen2.5-0.5B-Instruct), embedding (all-MiniLM-L6-v2 + FAISS) with Hugging Face

answer.txt

Вопрос: Кто создал Python?
Ответ: Гвидо ван Россум создал Python.
Источники: ['Python – высокоуровневый язык программирования общего назначения. Создан в 1991 году Гвидо ван Россумом.', 'FAISS – библиотека для эффективного поиска схожих векторов.', 'Эмбеддинги – это векторные представления текста, которые кодируют семантическое содержание в числовой форме.']
--------------------------------------------------
Вопрос: Что такое RAG?
Ответ: RAG (Retrieval-Augmented Generation) - это архитектурное решение, которое сочетает использование индексации для быстрого поиска и генерации текста.
Источники: ['RAG (Retrieval-Augmented Generation) – архитектурный подход, который объединяет поиск информации с генерацией текста.', 'Эмбеддинги – это векторные представления текста, которые кодируют семантическое содержание в числовой форме.', 'FAISS – библиотека для эффективного поиска схожих векторов.']
--------------------------------------------------
Вопрос: Как работают векторные базы данных?
Ответ: Векторные базы данных хранят данные в виде высокомерных векторов, которые позволяют выполнять семантический поиск.
Источники: ['Векторные базы данных хранят данные в виде высокомерных векторов и позволяют выполнять семантический поиск.', 'Эмбеддинги – это векторные представления текста, которые кодируют семантическое содержание в числовой форме.', 'FAISS – библиотека для эффективного поиска схожих векторов.']
--------------------------------------------------

min-rag.py

import numpy as np
import faiss
import torch
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, pipeline
from typing import List


class SimpleRAG:
    def __init__(self):
        # Инициализируем модель для создания эмбеддингов
        self.encoder = SentenceTransformer("all-MiniLM-L6-v2")
        # Хранилище для документов и индекс FAISS
        self.documents = []
        self.index = None

        # Загружаем инструктивную Qwen (causal, но умеет в чат)
        model_name = "Qwen/Qwen2.5-0.5B-Instruct"
        self.gen_tokenizer = AutoTokenizer.from_pretrained(model_name)
        # pipeline загружает модель внутри себя, используя переданный токенизатор
        self.generator = pipeline(
            "text-generation",
            model=model_name,
            tokenizer=self.gen_tokenizer,
            device=-1,
            dtype=torch.float32,
        )

    def add_documents(self, documents: List[str]):
        """Добавляем документы в базу знаний"""
        self.documents.extend(documents)
        # Генерируем эмбеддинги для всех документов
        embeddings = self.encoder.encode(self.documents)
        # Создаём FAISS индекс для быстрого поиска
        dimension = embeddings.shape[1]
        self.index = faiss.IndexFlatIP(
            dimension
        )  # Inner Product для косинусного сходства
        # Нормализуем эмбеддинги для корректной работы с cosine similarity
        faiss.normalize_L2(embeddings)
        self.index.add(embeddings.astype("float32"))

    def search(self, query: str, k: int = 3) -> List[str]:
        """Ищем наиболее релевантные документы"""
        # Векторизуем запрос
        query_embedding = self.encoder.encode([query])
        faiss.normalize_L2(query_embedding)
        # Выполняем поиск в FAISS
        scores, indices = self.index.search(query_embedding.astype("float32"), k)
        # Возвращаем найденные документы
        return [self.documents[idx] for idx in indices[0]]

    def generate_answer(self, query: str, context: List[str]) -> str:
        """Генерируем ответ при помощи rut5-base-multitask"""
        # Формируем структурированный контекст с нумерацией документов
        context_text = "\n".join(
            [f"Документ {i+1}: {doc}" for i, doc in enumerate(context)]
        )

        messages = [
            {
                "role": "system",
                "content": "Ты — полезный ассистент. Отвечай на вопрос, используя только предоставленный контекст. Перефразируй информацию полным предложением на русском языке.",
            },
            {
                "role": "user",
                "content": f"Контекст:\n{context_text}\n\nВопрос: {query}",
            },
        ]

        out = self.generator(
            messages,
            max_new_tokens=100,
            do_sample=False,
            pad_token_id=self.gen_tokenizer.pad_token_id,
            eos_token_id=self.gen_tokenizer.eos_token_id,
            return_full_text=False,
        )
        return out[0]["generated_text"].strip()

    def ask(self, query: str) -> dict:
        """Основная функция: поиск + генерация ответа"""
        # Находим релевантные документы
        relevant_docs = self.search(query)
        # Генерируем ответ
        answer = self.generate_answer(query, relevant_docs)
        return {"query": query, "answer": answer, "sources": relevant_docs}


# Создаём экземпляр RAG-системы
rag = SimpleRAG()

# Добавляем документы
documents = [
    "Python – высокоуровневый язык программирования общего назначения. Создан в 1991 году Гвидо ван Россумом.",
    "RAG (Retrieval-Augmented Generation) – архитектурный подход, который объединяет поиск информации с генерацией текста.",
    "Векторные базы данных хранят данные в виде высокомерных векторов и позволяют выполнять семантический поиск.",
    "FAISS – библиотека для эффективного поиска схожих векторов.",
    "Эмбеддинги – это векторные представления текста, которые кодируют семантическое содержание в числовой форме.",
]
rag.add_documents(documents)

# Задаём вопросы
for question in [
    "Кто создал Python?",
    "Что такое RAG?",
    "Как работают векторные базы данных?",
]:
    result = rag.ask(question)
    print(f"Вопрос: {result['query']}")
    print(f"Ответ: {result['answer']}")
    print(f"Источники: {result['sources']}")
    print("-" * 50)