Spaces:
Build error
Build error
File size: 3,873 Bytes
5f0bb70 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
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
)
|