🤖 Дедубликация: как OpenAI и FastAPI спасут Habr от дублей

doc А = 'Цена нефти бренд выросла из-за опасений возможности столкновения между США и Ираном'

doc B = 'Стоимость североморской нефти растет на ожидании обострения конфликта Пентагона и Ирана’

В данном случае мера Жаккара будет около 0.1

from sentence_transformers import SentenceTransformer
# импортируем мультиязычную модель
model = SentenceTransformer('distiluse-base-multilingual-cased-v2')
sentence = 'Мама мыла раму'
# получаем эмбеддинг
embedding = model.encode(sentence)

sentence_1 = 'Мама мыла раму' * 64
sentence_2 = 'Мама мыла раму' * 128
embeddings = model.encode([sentence_1, sentence_2])
print((embeddings[0] == embeddings[1]).all())
True

conda create --name open_ai
conda activate open_ai

pip install openai tiktoken juputer scipy pandas fastapi uvicorn

import tiktoken
import openai
from scipy.spatial.distance import cosine
# Устанавливаем значение ключа
openai.api_key = 'sk-7FUL3uCuviG7F...Kc08xiSoFidB9Us'
# Примеры предложений
sentence_1 = 'мама мыла раму'
sentence_2 = 'мама моет окно'
# Версия модели(это наиболее современная)
model = 'text-embedding-ada-002'
# Создаем эмбеддинги
embedding_1 = openai.Embedding.create( input=sentence_1, model=model)
embedding_2 = openai.Embedding.create( input=sentence_2, model=model)
# Смотрим на близость, в данном случае считается как 1-cos
distance = cosine(embedding_1['data'][0]['embedding'], embedding_2['data'][0]['embedding'])
print(distance)
0.11193540954120007

sentence_1 = '''
В январе Генеральная прокуратура Германии в рамках расследования диверсии на газопроводах «Северный поток» проводила обыск
на подозрительном судне, которое предположительно перевозило взрывчатку, пишет Die Welt со ссылкой на заявление ведомства.
Подробностей, кто и зачем перевозил ее, ведомство не приводит. В ответ на запрос ТАСС в Генпрокуратуре указали, что продолжают
анализировать улики и не могут говорить о причастности к подрыву какого-либо государства. Судно, по данным ведомства,
было арендовано у компании из ФРГ, но ее сотрудники вне подозрений.
Ранее газета Die Zeit сообщила, что немецкие следственные органы идентифицировали яхту,
которая, вероятно, использовалась при совершении взрыва на «Северных потоках» 26 сентября прошлого года.
'''
sentence_2 = '''
Власти ФРГ инициировали обыск корабля, предположительно имеющего отношение к диверсиям на газопроводах «Северный поток»
и «Северный поток – 2». Об этом 8 марта сообщает агентство DPA со ссылкой на Генеральную прокуратуру в Карлсруэ.
Согласно информации ведомства, судно могло быть использовано для перевозки взрывчатых веществ.
«С 18 по 20 января 2023 года федеральная прокуратура обыскала судно в связи с подозрительной его арендой.
Есть подозрение, что рассматриваемое судно могло использоваться для перевозки взрывных устройств,
взорвавшихся 26 сентября 2022 года на газопроводах «Северный поток» 1 и 2 в Балтийском море»,
— сообщили в прокуратуре агентству «РИА Новости».
Отмечается, что на данный момент изъят ряд вещдоков, проводится их проверка. Личности преступников
и их мотивы являются предметом продолжающихся расследований.
«Достоверных доказательств по этому поводу, особенно по вопросу государственного контроля,
в настоящее время нет. Никаких подозрений в отношении сотрудников немецкой компании,
сдавшей судно в аренду, нет», — заключили в прокуратуре.

'''
# Энкодер для подсчета количества токенов
encoding = tiktoken.get_encoding('cl100k_base')

num_tokens_1 = len(encoding.encode(sentence_1))
num_tokens_2 = len(encoding.encode(sentence_2))

embedding_1 = openai.Embedding.create(
  input=sentence_1,
  model=model
)

embedding_2 = openai.Embedding.create(
  input=sentence_2,
  model=model
)

distance = cosine(
	embedding_1['data'][0]['embedding'],
	embedding_2['data'][0]['embedding']
)

print(
    f'sentence_1 num tokens: {num_tokens_1}\\n'
    f'sentence_2 num tokens: {num_tokens_2}\\n'
    f'cos distance: {distance}'
)

sentence_1 num tokens: 347
sentence_2 num tokens: 492
cos distance: 0.06688715737425832

api_input = {
	  "ITEMS": [
	    doc1,
	    doc2,
	    doc3
	  ]
	}
	
api_output = [
    {
        "estimation": 1,
        "score": 0.3,
        "duplicates": []
    },
    {
        "estimation": 0,
        "score": 0.03,
        "duplicates": [original1, original2]
    },
    {
        "estimation": 1,
        "score": 0.41,
        "duplicates": []
    },
]

import pandas as pd
import pickle

# Создаем итератор для побатчевой загрузки, возможно потребуется установка стандарта zstd(pip install zstandard)
habr = pd.read_json(
	'habr.jsonl.zst',
	lines=True,
	chunksize=10**4,
	compression='zstd'
)
# Забираем первый батч
habr_chunk = next(habr)

print(habr_chunk.shape)
(10000, 22)

def get_embeddings(  
        texts: list,  
        model: str = 'text-embedding-ada-002',  
        chunksize: int = 100,  
        max_tokens: int = 8192,  
        cut_coef: float = 0.8  
) -> list:  
    """  
    
Функция принимает список текстов и возвращает список эмбеддингов, используя    модель от OpenAI. Список разбивается на куски размером chunksize, а затем функция вычисляет количество    токенов каждого текста, используя энкодер cl100k_base. Если количество токенов превышает max_tokens,текст укорачивается на cut_coef.
  
Аргументы:
  
    df (pd.DataFrame): pandas DataFrame, содержащий текстовые данные
		model (str по умолчанию text-embedding-ada-002): модель векторизации
    chunksize(int по умолчанию 100): размер кусков, на которые разбиваются входные данные
    max_tokens (int по умолчанию 8192): максимальное количество токенов, разрешенное на один текст
    cut_coef (float по умолчанию 0.8): коэффициент, используемый для определения количества текста для обрезки, если оно превышает max_tokens
  
Возвращает:

    список эмбеддингов  
    
    """  
    embeddings_all = []  
    encoding = tiktoken.get_encoding('cl100k_base')  
  
    @retry(wait=wait_random_exponential(min=1, max=10), stop=stop_after_attempt(5))  # обработчик запросов  
    def api_request(texts, model):  
        return openai.Embedding.create(input=texts, model=model)  
  
    for i in range(0, len(texts), chunksize):  
        processed_texts = []  
        texts_chunk = texts[i:i + chunksize]  
        for text in texts_chunk:  
            if not text:  
                text = 'default'  
            num_tokens = len(encoding.encode(text))  
            if num_tokens > max_tokens:  
                len_text = len(text)  
                cut_trh = int(max_tokens / num_tokens * len_text * cut_coef)  
                text = text[:cut_trh]  
            processed_texts.append(text)  
  
        embeddings = api_request(processed_texts, model)  
  
        for embedding in embeddings['data']:  
            embeddings_all.append(embedding['embedding'])  
        if (i % 1000 == 0) & (i != 0):  
            print(f'processed {i} texts')  
  
    return embeddings_all

embeddings_all = get_embeddings(habr_chunk.text_markdown.values.tolist())

habr_chunk['embedding'] = embeddings_all

with open('habr_chunk.pickle', 'wb') as f:
    pickle.dump(habr_chunk[['text_markdown', 'embedding']], f

def find_duplicates(candidates: list, df: pd.DataFrame, n_similar: int = 5, thrh: float = 0.1) -> list:
    '''
	Функция находит дубликаты текстовых фрагментов для  списка candidates на основе сходства эмбеддингов, вычисленных с помощью функции get_embeddings, и базы данных в виде pandas.DataFrame. Функция возвращает список, содержащий оценку, близость и список похожих фрагментов для каждого текстового фрагмента из candidates.

 Аргументы:

    candidates (list): список текстовых фрагментов для поиска дубликатов
    df (pd.DataFrame): база данных в формате pandas.DataFrame, содержащая эмбеддинги и текстовые фрагменты для сравнения с candidates
    n_similar (int, по умолчанию 5): количество похожих фрагментов, которые будут возвращены для каждого текста candidates
    thrh (float, по умолчанию 0.1): пороговое значение, используемое для определения того, что два эмбеддинга достаточно похожи.
    
 Возвращает:

    список словарей с оценкой, косинусной близостью и списком дубликатов.
    '''
    embeddings = get_embeddings(candidates)
    processed_candidates = []
    for embedding in embeddings:
        checked_candidat = {
            "estimation": 0,
            "score": None,
            "duplicates": []
        }
        distances = np.array(  
            [cosine(embedding, e) for e in df.embedding.values]  
        )
        under_treshold = distances < thrh
        if under_treshold.sum() == 0:  
            processed_candidates.append(checked_candidat)  
        else: 
            checked_candidat["estimation"] = 1
            checked_candidat["score"] = min(distances)
            checked_candidat["duplicates"] = df.iloc[
                under_treshold
            ].text_markdown.values[:n_similar].tolist()
            
        processed_candidates.append(checked_candidat)
        
    return processed_candidates

mkdir utils
touch utils/__init__.py utils/search.py app.py

import os  
import pickle  
  
from fastapi import FastAPI  
from pydantic import BaseModel  
  
import openai  
  
from utils.search import find_duplicates  
  
# Устанавливаем значение ключа(так лучше не хранить!!!)
openai.api_key = 'sk-7F...idB9Us'  
  
# Название файла с уникальными постами  
habr_masters_path = 'habr_chunk.pickle'  
  
# Создаем экземпляр fastapi  
app = FastAPI(title='Deduplication App')  
  
# Для валидации данных создаем модель данных 
class Candidat(BaseModel):  
    ITEM: list  
   
# Подгружаем датафрейм с постами  
with open(os.path.join(os.path.dirname(__file__), habr_masters_path), 'rb') as f:  
    df = pickle.load(f)  

# Обработчик для входящих сообщений
@app.post('/predict')  
async def make_predictions(request: Candidat):  
    if not request.ITEM:  
        request.ITEM.append('dafault')  
    processed_candidates = find_duplicates(request.ITEM, df)  
    return processed_candidates

uvicorn app:app --reload

INFO:     Will watch for changes in these directories: ['/Users/m.konakov/Desktop/search_duplicates']
INFO:     Uvicorn running on <http://127.0.0.1:8000> (Press CTRL+C to quit)
INFO:     Started reloader process [24198] using StatReload
INFO:     Started server process [24200]
INFO:     Waiting for application startup.
INFO:     Application startup complete.