143 lines
5.6 KiB
Python
143 lines
5.6 KiB
Python
import pandas as pd
|
||
import numpy as np
|
||
import joblib
|
||
import os
|
||
from sklearn.model_selection import train_test_split
|
||
from sklearn.feature_extraction.text import TfidfVectorizer
|
||
from sklearn.linear_model import LogisticRegression
|
||
from sklearn.pipeline import Pipeline, FeatureUnion
|
||
from sklearn.preprocessing import StandardScaler
|
||
from sklearn.base import BaseEstimator, TransformerMixin
|
||
from sklearn.metrics import accuracy_score, classification_report
|
||
import rich # <!-- дебаг через rich - моя guilty pleasure, очень уж люблю на красивые выводы смотреть
|
||
from rich import print as rpint
|
||
from rich.console import Console
|
||
from rich import box
|
||
from rich.table import Table
|
||
from rich.markdown import Markdown
|
||
|
||
console = Console()
|
||
|
||
|
||
# инициализируем специальные классы, чтобы раскидать данные по категориям
|
||
|
||
# Для преобразования TF-IDF в вектора
|
||
class TextExtractor(BaseEstimator, TransformerMixin):
|
||
def fit(self, X, y=None): return self
|
||
def transform(self, X):
|
||
return X['full_text'].fillna('')
|
||
|
||
# Для StandardScaler
|
||
class NumberExtractor(BaseEstimator, TransformerMixin):
|
||
def fit(self, X, y=None): return self
|
||
def transform(self, X):
|
||
return X[['amount']].fillna(0)
|
||
|
||
|
||
|
||
|
||
def train_model():
|
||
console.log("[yellow]Грузим данные из data...[/yellow]")
|
||
try:
|
||
tx = pd.read_csv('data/transactions.csv')
|
||
terminals = pd.read_csv('data/terminals.csv')
|
||
receipts = pd.read_csv('data/receipts.csv')
|
||
except FileNotFoundError as e:
|
||
console.log(f"Файлы для обучения не найдены :( \n {e}", style="white on red")
|
||
return
|
||
|
||
|
||
console.log("[yellow]Предобрабатываем данные...[/yellow]")
|
||
# Приклеиваеем вместе имена товаров
|
||
receipts_agg = receipts.groupby('transaction_id')['item_name'].apply(
|
||
lambda x: ' '.join(str(i) for i in x)
|
||
).reset_index()
|
||
|
||
# Делаем один большой датафрейм с которым будем работать
|
||
df = tx.merge(terminals[['terminal_id', 'terminal_name', 'terminal_description']], on='terminal_id', how='left')
|
||
df = df.merge(receipts_agg, on='transaction_id', how='left')
|
||
|
||
# Делаем текстовое поле для TF-IDF
|
||
df['full_text'] = (
|
||
df['terminal_name'].astype(str) + " " +
|
||
df['terminal_description'].astype(str) + " " + # <!-- изначально я пробовал клеить id транзакции, однако модель слишком на ней зацикливалась
|
||
df['item_name'].astype(str)
|
||
).str.lower()
|
||
|
||
X = df[['full_text', 'amount']]
|
||
y = df['true_mcc']
|
||
|
||
# Пайплайн обучения
|
||
|
||
pipeline = Pipeline([
|
||
('features', FeatureUnion([
|
||
# Ветка для слов
|
||
('text_branch', Pipeline([
|
||
('extract', TextExtractor()),
|
||
('tfidf_union', FeatureUnion([
|
||
# Векторизуем слова и удаляем лишние слова без смысла
|
||
('word', TfidfVectorizer(
|
||
ngram_range=(1, 2),
|
||
analyzer='word',
|
||
stop_words='english',
|
||
max_features=5000
|
||
)),
|
||
# Фиксим очепятки
|
||
('char', TfidfVectorizer(
|
||
ngram_range=(2, 5),
|
||
analyzer='char_wb',
|
||
max_features=10000
|
||
))
|
||
]))
|
||
])),
|
||
# Ветка для чисел
|
||
('numeric_branch', Pipeline([
|
||
('extract', NumberExtractor()),
|
||
('scaler', StandardScaler())
|
||
]))
|
||
])),
|
||
# Для классификации юзаем логрег
|
||
('clf', LogisticRegression(C=1.0, max_iter=1000)) # <!-- были разные коэффиценты, в итоге оставил что-то средненькое
|
||
])
|
||
|
||
# Валидация
|
||
|
||
console.log("[yellow]Оцениваем качество на валидационной выборке...[/yellow]")
|
||
X_train, X_test, y_train, y_test = train_test_split(
|
||
X, y, test_size=0.2, random_state=42, stratify=y
|
||
)
|
||
|
||
|
||
|
||
pipeline.fit(X_train, y_train)
|
||
y_pred = pipeline.predict(X_test)
|
||
probs = pipeline.predict_proba(X_test)
|
||
|
||
acc = accuracy_score(y_test, y_pred)
|
||
conf = np.mean(np.max(probs, axis=1))
|
||
|
||
table = Table(box=box.ROUNDED, title="Отчёт")
|
||
|
||
|
||
table.add_column("Метрика", justify="center", style="yellow")
|
||
table.add_column("Значение", justify="center", style="yellow")
|
||
|
||
table.add_row("Accuracy", f"{acc:.4f}")
|
||
table.add_row("Avg Confidence", f"{conf:.4f}")
|
||
console.print(table, justify="center")
|
||
|
||
console.print("[yellow]Репорт по классам[/yellow]", justify="center")
|
||
console.print(classification_report(y_test, y_pred), justify="center")
|
||
|
||
# Метрики норм, учимся на всех данных и сохранем данные
|
||
|
||
console.log("[yellow]Учимся на всем, что есть...[/yellow]")
|
||
pipeline.fit(X, y)
|
||
|
||
os.makedirs('solution/model', exist_ok=True)
|
||
joblib.dump(pipeline, 'solution/model/mcc_model.pkl')
|
||
console.log(Markdown("Сохранили модель в **solution/model/mcc_model.pkl**"))
|
||
|
||
with console.status("Учим модель..."):
|
||
train_model()
|
||
console.print(Markdown("*Модель готова :)*")) |