106 lines
2.5 KiB
Go
106 lines
2.5 KiB
Go
package dedup
|
|
|
|
import (
|
|
"context"
|
|
"sort"
|
|
|
|
"github.com/emil/deepres/internal/research/embed"
|
|
"github.com/emil/deepres/internal/research/search"
|
|
)
|
|
|
|
// DedupResult содержит результаты после дедупликации
|
|
type DedupResult struct {
|
|
Items []DedupItem
|
|
RawTexts []string
|
|
}
|
|
|
|
type DedupItem struct {
|
|
Title string
|
|
URL string
|
|
Content string
|
|
ClusterID int
|
|
}
|
|
|
|
// Deduplicate принимает результаты поиска, эмбеддит, кластеризует и дедуплицирует
|
|
func Deduplicate(ctx context.Context, embedCli *embed.EmbedClient, results []search.SearXNGResult, threshold float32) (*DedupResult, error) {
|
|
if len(results) == 0 {
|
|
return &DedupResult{}, nil
|
|
}
|
|
|
|
// Подготавливаем тексты для эмбеддинга
|
|
texts := make([]string, len(results))
|
|
for i, r := range results {
|
|
content := r.Title + ". " + r.Content
|
|
if len([]rune(content)) > 512 {
|
|
content = string([]rune(content)[:512])
|
|
}
|
|
texts[i] = content
|
|
}
|
|
|
|
// Эмбеддинги
|
|
embeddings, err := embedCli.Embed(ctx, texts)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Простая дедупликация: если косинусная близость > threshold — это дубликат
|
|
n := len(results)
|
|
clusters := make([]int, n)
|
|
clusterID := 0
|
|
for i := range clusters {
|
|
clusters[i] = -1
|
|
}
|
|
|
|
for i := 0; i < n; i++ {
|
|
if clusters[i] != -1 {
|
|
continue
|
|
}
|
|
clusters[i] = clusterID
|
|
for j := i + 1; j < n; j++ {
|
|
if clusters[j] != -1 {
|
|
continue
|
|
}
|
|
sim := embed.CosineSimilarity(embeddings[i], embeddings[j])
|
|
if sim >= threshold {
|
|
clusters[j] = clusterID
|
|
}
|
|
}
|
|
clusterID++
|
|
}
|
|
|
|
// Собираем результат: по одному лучшему элементу из каждого кластера
|
|
unique := make(map[int]*DedupItem)
|
|
rawTexts := make([]string, 0)
|
|
|
|
// Сначала собираем все
|
|
for i, r := range results {
|
|
key := clusters[i]
|
|
if _, exists := unique[key]; !exists {
|
|
unique[key] = &DedupItem{
|
|
Title: r.Title,
|
|
URL: r.URL,
|
|
Content: r.Content,
|
|
}
|
|
}
|
|
}
|
|
|
|
// Сортируем по clusterID
|
|
keys := make([]int, 0, len(unique))
|
|
for k := range unique {
|
|
keys = append(keys, k)
|
|
}
|
|
sort.Ints(keys)
|
|
|
|
var items []DedupItem
|
|
for _, k := range keys {
|
|
item := unique[k]
|
|
items = append(items, *item)
|
|
rawTexts = append(rawTexts, item.Title+". "+item.Content)
|
|
}
|
|
|
|
return &DedupResult{
|
|
Items: items,
|
|
RawTexts: rawTexts,
|
|
}, nil
|
|
}
|