Помоги мне, пожалуйста, как-то добавить word2vec в мой код. Моя цель - улучшение f1-score
Вот мой код:
# Импортирование необходимых библиотек
import random
from typing import Dict, List, Union

from datasets import load_dataset
from fuzzywuzzy import fuzz, process
from sklearn.metrics import classification_report
from tqdm.notebook import tqdm
from transformers import pipeline
import torch

# Ваша функция для подготовки сообщений осталась без изменений
def prepare_message_for_llm(text: Union[str, List[str]], categories: Dict[str, str]) -> Dict[str, Union[List[Dict[str, str]], List[List[Dict[str, str]]]]]:
    if len(categories) < 2:
        raise RuntimeError(f'The category list is too small! Expected 2 or more categories, got {len(categories)} ones.')
    categories_ = sorted(list(categories.keys()))
    categories_as_string = ', '.join(categories_[:-1]) + ' и ' + categories_[-1]
    if isinstance(text, str):
        prompt = f'Прочтите, пожалуйста, следующий текст и определите, какая тема из известного ' \
                 f'списка тем наиболее представлена в следующем тексте. ' \
                 f'В качестве ответа напишите только название темы из списка, больше ничего.\n' \
                 f'Список тем: {categories_as_string}.\n'
        for cur in categories_:
            prompt += f'Текст: {" ".join(categories[cur].split())}\nВаш ответ: {cur}\n'
        prompt += f'Текст: {" ".join(text.split())}\nВаш ответ: '
        messages = [
            {
                'role': 'system',
                'content': 'Вы - полезный помощник, умеющий читать тексты на русском языке, глубоко понимать их и анализировать.'
            },
            {
                'role': 'user',
                'content': prompt
            }
        ]
    else:
        messages = []
        for it in text:
            prompt = f'Прочтите, пожалуйста, следующий текст и определите, какая тема из известного ' \
                     f'списка тем наиболее представлена в следующем тексте. ' \
                     f'В качестве ответа напишите только название темы из списка, больше ничего.\n' \
                     f'Список тем: {categories_as_string}.\n'
            for cur in categories_:
                prompt += f'Текст: {" ".join(categories[cur].split())}\nВаш ответ: {cur}\n'
            prompt += f'Текст: {" ".join(text.split())}\nВаш ответ: '
            messages.append([
                {
                    'role': 'system',
                    'content': 'Вы - полезный помощник, умеющий читать тексты на русском языке, глубоко понимать их и анализировать.'
                },
                {
                    'role': 'user',
                    'content': prompt
                }
            ])
    return {'message_for_llm': messages}

# Инициализация pipeline для zero-shot классификации
llm_pipeline = pipeline(
    'zero-shot-classification',
    model='distilbert-base-multilingual-cased',  # Multilingual поддерживает русский язык
    device=0 if torch.cuda.is_available() else -1  # Использовать GPU, если доступен
)

# Загрузка данных
DATASET_NAME = 'Davlan/sib200'
DATASET_LANGUAGE = 'rus_Cyrl'
train_set = load_dataset(DATASET_NAME, DATASET_LANGUAGE, split='train')
validation_set = load_dataset(DATASET_NAME, DATASET_LANGUAGE, split='validation')
test_set = load_dataset(DATASET_NAME, DATASET_LANGUAGE, split='test')

# Определение категорий
list_of_categories = sorted(list(
    set(train_set['category']) | set(validation_set['category']) | set(test_set['category'])
))
print(f'Categories for classification are: {list_of_categories}')

# Подготовка примеров для каждой категории
examples_by_categories = {}
for current_category in list_of_categories:
    examples_by_categories[current_category] = random.choice(
        train_set.filter(lambda it: it['category'] == current_category)['text']
    )

# Подготовка данных для LLM
validation_set_for_llm = validation_set.map(lambda it: prepare_message_for_llm(it['text'], examples_by_categories))
test_set_for_llm = test_set.map(lambda it: prepare_message_for_llm(it['text'], examples_by_categories))

# Уменьшение размера выборки
sample_size = 10  # Пример: используем 10 примеров
validation_sample = validation_set_for_llm.select(range(sample_size))

# Генерация предсказаний
y_pred = list(map(
    lambda messages: llm_pipeline(
        messages[1]['content'],  # Доступ к содержимому текста сообщения
        candidate_labels=list_of_categories,
        hypothesis_template="Этот текст связан с категорией {}."
    )['labels'][0],  # Получение первой прогнозируемой метки
    tqdm(validation_sample['message_for_llm'])
))

# Истинные метки
y_true = validation_set.select(range(sample_size))['category']

# Печать отчета о классификации
print(classification_report(y_true=y_true, y_pred=y_pred))