finetune.py

from transformers import AutoTokenizer, AutoModel
import torch
import faiss
import numpy as np
import gradio as gr


# Загружаем токенизатор и модель
MODEL_NAME = "distilbert-base-multilingual-cased"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME)

# Тестовая фраза
text = "Привет, как дела?"

# Токенизация
tokens = tokenizer(text, return_tensors="pt", padding=True, truncation=True)

# Генерация эмбеддингов
with torch.no_grad():
    output = model(**tokens)

# Вывод размера эмбеддингов
embedding = output.last_hidden_state[:, 0, :]  # Берём CLS-токен
print("Размер эмбеддинга:", embedding.shape)


# Размерность эмбеддинга (768, как у DistilBERT)
D = 768  

# Создаём FAISS-индекс (L2 - евклидово расстояние)
index = faiss.IndexFlatL2(D)

# Проверяем, пуст ли индекс
print("Индекс создан. Количество векторов:", index.ntotal)

import json

# Загрузка данных из файла
with open('data.json', 'r', encoding='utf-8') as f:
    documents = json.load(f)

# Проверим данные
print(documents[:5])  # Печатаем первые 5 элементов


# Кодируем тексты в эмбеддинги
doc_embeddings = []
for doc in documents:
    tokens = tokenizer(doc, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        embedding = model(**tokens).last_hidden_state[:, 0, :].numpy()  # Берём CLS-токен
    doc_embeddings.append(embedding)

# Превращаем список в numpy-массив
doc_embeddings = np.vstack(doc_embeddings)

# Добавляем в FAISS
index.add(doc_embeddings)

# Проверяем, сколько векторов в базе
print("Векторов в индексе:", index.ntotal)

def search_query(query, index, top_k=3):
    # Токенизация запроса
    tokens = tokenizer(query, return_tensors="pt", padding=True, truncation=True)
    
    # Генерация эмбеддинга для запроса
    with torch.no_grad():
        query_embedding = model(**tokens).last_hidden_state[:, 0, :].numpy()
    
    # Поиск ближайших векторов
    D, I = index.search(query_embedding, top_k)
    
    return D, I

# Пример запроса
query = "веб-сервис"

# Получаем результаты поиска
distances, indices = search_query(query, index)

# Печатаем результаты
print("Результаты поиска:")
for dist, idx in zip(distances[0], indices[0]):
    print(f"Текст: {documents[idx]} | Расстояние: {dist}")


def search_in_faiss(query: str):
    # Преобразуем запрос в эмбеддинг
    tokens = tokenizer(query, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        query_embedding = model(**tokens).last_hidden_state[:, 0, :].numpy()

    # Находим ближайшие векторы в FAISS
    distances, indices = index.search(query_embedding, k=5)

    # Формируем результаты
    results = [documents[idx] for idx in indices[0]]
    return results

iface = gr.Interface(
    fn=search_in_faiss,  # Функция для поиска в FAISS
    inputs="text",  # Ввод текста
    outputs="text",  # Вывод результата (тексты документов)
    title="Поиск по проектам",
    description="Введите запрос, и система найдет наиболее подходящие ответы из базы знаний."
)

# Запуск интерфейса
iface.launch(
    share=True
)