Ниже — готовый README.md для вашего main.py. Он описывает назначение файла, структуру, требования, запуск, а также функциональность.
Sentiment Analysis Pipeline (RuBERT Tiny)
Этот проект реализует полный конвейер обработки текстов и обучения модели для классификации тональности русскоязычных текстов. Используется модель cointegrated/rubert-tiny-sentiment-balanced, а также квантование через Intel Neural Compressor для оптимизации до INT8.
Возможности
- Очистка и лемматизация русских текстов с использованием pymorphy3 и NLTK
- Удаление стоп-слов
- Токенизация и подготовка датасета для HuggingFace
- Обучение модели RuBERT Tiny на ваших данных
- Оценка метрик: Accuracy, F1, Precision, Recall
- Квантование обученной модели до INT8
- Инференс для одиночных текстов или DataFrame
- Возможность легко интегрировать обученную модель
##Структура проекта
Основные функции из main.py:
###load_data()
Загружает файлы train.csv и test.csv.
###preprocess_text(text)
Нормализует текст:
- нижний регистр
- удаление пунктуации
- токенизация
- лемматизация
- фильтрация стоп-слов
###clean_df(df)
Добавляет столбец clean_text.
###train_model(train_df)
Обучает RuBERT Tiny:
- токенизация
- разбиение на train/val
- обучение через Trainer
- сохранение модели в rubert_tiny_sentiment_model_fp32
###quantize_model()
Квантование FP32 → INT8 Сохраняет модель в rubert_tiny_sentiment_model_int8.
###run_inference(model, tokenizer)
Инференс для одного текста.
predict_dataframe(model, tokenizer, df)
Инференс для DataFrame.
###main()
Полный цикл:
- загрузка данных
- очистка
- обучение
- квантование
##Требования
Установите зависимости:
pip install pandas numpy nltk torch evaluate transformers datasets pymorphy3 optimum[neural-compressor] Перед первым запуском NLTK докачает токенизаторы автоматически.
##Запуск обучения
Просто запустите:
python main.py После обучения появятся директории:
rubert_tiny_sentiment_model_fp32/ rubert_tiny_sentiment_model_int8/
##Использование инференса
Одиночный текст
from main import run_inference out = run_inference(model, tokenizer)("мне очень понравился товар!") print(out)
{'label': 'positive', 'confidence': 0.94}
DataFrame
df = predict_dataframe(model, tokenizer, df) print(df)
Требуемый формат train.csv / test.csv:
| text | label |
|---|---|
| "Этот фильм отличный" | 1 |
| "Мне не понравилось" | 2 |
Где:
- 0 — neutral
- 1 — positive
- 2 — negative
##Лицензия
Проект создан для демонстрации работы с RuBERT и Intel Neural Compressor. Все используемые библиотеки распространяются по своим лицензиям.