!pip install datasets

!pip install transformers

!pip install accelerate

!pip install fuzzywuzzy

!pip install python-Levenshtein

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

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}

llm_pipeline = pipeline(model='Qwen/Qwen2-7B-Instruct', device_map='auto', torch_dtype='auto')

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}')

print(validation_set)

examples_by_categories = dict()
for current_category in list_of_categories:
    examples_by_categories[current_category] = random.choice(
        train_set.filter(lambda it: it['category'] == current_category)['text']
    )
    print(f'Category: {current_category}\n')
    print(f'Random text: {examples_by_categories[current_category]}\n\n')

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))

print(validation_set_for_llm)

print(validation_set['text'][0])

print(validation_set_for_llm['message_for_llm'][0])

y_pred = list(map(
    lambda x: llm_pipeline(x, max_new_tokens=10)[0]['generated_text'],
    tqdm(validation_set_for_llm['message_for_llm'])
))
y_true = validation_set['category']

print(y_pred[0])

print(classification_report(y_true=y_true, y_pred=[x[-1]['content'] for x in y_pred]))

y_pred_with_normalization = list(map(
    lambda it: process.extractOne(it[-1]['content'], list_of_categories, scorer=fuzz.token_sort_ratio)[0],
    y_pred
))

print(y_pred_with_normalization[0])

print(classification_report(y_true=y_true, y_pred=y_pred_with_normalization))

y_pred = list(map(
    lambda x: llm_pipeline(x, max_new_tokens=10)[0]['generated_text'],
    tqdm(test_set_for_llm['message_for_llm'])
))
y_true = test_set['category']

print(classification_report(y_true=y_true, y_pred=[x[-1]['content'] for x in y_pred]))

y_pred_with_normalization = list(map(
    lambda it: process.extractOne(it[-1]['content'], list_of_categories, scorer=fuzz.token_sort_ratio)[0],
    y_pred
))

print(classification_report(y_true=y_true, y_pred=y_pred_with_normalization))