ml_bert_test.py - master - maratgaliulin/landcode_classifier

ml_bert_test.py
198 строк7 KB
Nutten Chestnut
Внёс правки Артёма в код
09 фев 2026, 04:34
09 фев 2026, 04:340dcabef
import pandas as pd
import os
import torch
import time

from torch.optim import AdamW
from torch.utils.data import DataLoader
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from transformers import AutoTokenizer, AutoModelForSequenceClassification

from methods.utils import save_full_model, return_dataframes_from_pickle_or_from_excel, download_model_and_components_from_cloud
from methods.classes.LandUsePredictor import LandUsePredictor
from methods.classes.BertLandUseDataset import BertLandUseDataset
from methods.classes.BertWithNumeric import BertWithNumeric


pickle_df_ids_dir = 'files/pickle/Коды ЗУ.pkl'
pickle_df_intermed_dir = 'files/pickle/тест ИИ.pkl'

excel_df_ids_dir = 'files/Коды ЗУ.xlsx'
excel_df_intermed_dir = 'files/тест ИИ.xlsx'

pickle_df_dir = 'files/pickle/df.pkl'

pickle_tfidf_dir = 'files/pickle/dfidf_df.pkl'

merged_df_dir = 'files/pickle/df_final.pkl'

text_columns = [
    'Вид использования участка по документу (САМЫЙ ГЛАВНЫЙ АТРИБУТ - Приоритет 1)',
    'Разрешенное использование (текстовое описание) - Приоритет 2',
    'Вид земельного участка',
    'Категория земель'
    ]

model_components_directory = 'files/pickle/ml_models/model_components.pkl'
model_directory = 'files/pickle/ml_models/rubert_landuse_model_without_area.pth'


model_components_public_key = 'https://disk.yandex.ru/d/YR1lL2k85Xkaqg'
model_public_key = 'https://disk.yandex.ru/d/0pvpfxulPy03xw'


zu_codes_public_key = 'https://disk.yandex.ru/i/9pvmRyxJGJUSRg'
ai_test_public_key = 'https://disk.yandex.ru/i/fHUNeZlMzEASBw'


os_path = 'files/pickle/ml_models/rubert_landuse_model_priority_first.pth'
components_path = 'files/pickle/ml_models/model_components_priority_first.pkl'


model_name = 'DeepPavlov/rubert-base-cased'

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

if(not os.path.exists('files/')):
    print('Необходимые директории отсутствуют. Создаю нужные директории.')
    os.mkdir('files/')
    os.mkdir('files/pickle/')
    os.mkdir('files/pickle/ml_models/')
    print('Нужные директории созданы.')
    time.sleep(3)   

else: print('Нужные директории уже имеются.')

if(not os.path.isfile(model_components_directory) or not os.path.isfile(model_directory)):
    
    print('Начинаю загрузку вспомогательных компонентов модели...')
    download_model_and_components_from_cloud(public_key=model_components_public_key, file_directory=model_components_directory)

    print('Вспомогательные компоненты модели успешно загружены. Начинаю загрузку модели...')
    download_model_and_components_from_cloud(public_key=model_public_key, file_directory=model_directory)

    print('Модель успешно загружена.')

else: print('Модели загружены.')



if(not os.path.isfile(excel_df_ids_dir) or not os.path.isfile(excel_df_intermed_dir)):
    
    print('Начинаю загрузку файлов Excel для обучения и тестирования...')
    download_model_and_components_from_cloud(public_key=zu_codes_public_key, file_directory=excel_df_ids_dir)

    download_model_and_components_from_cloud(public_key=ai_test_public_key, file_directory=excel_df_intermed_dir)


else: print('Файлы Excel загружены.')



if(not os.path.isfile(os_path)):
    
    df_ids, df = return_dataframes_from_pickle_or_from_excel(
        pickle_df_ids_dir=pickle_df_ids_dir, 
        pickle_df_intermed_dir=pickle_df_intermed_dir, 
        excel_df_ids_dir=excel_df_ids_dir, 
        excel_df_intermed_dir=excel_df_intermed_dir
        )

    # Prepare target
    label_encoder = LabelEncoder()
    label_encoder.fit(df_ids['Код расчета вида использования'])
    df_ids['label'] = label_encoder.transform(df_ids['Код расчета вида использования'])
    df['label'] = label_encoder.transform(df['Код расчёта вида использования (ГБУ)'])

    train_texts, val_texts, train_labels, val_labels = train_test_split(
        df['combined_text'].values,
        df['label'].values,
        test_size=0.2,
        random_state=42
    )

    train_texts1 = df_ids['combined_text'].values
    train_labels1 = df_ids['label'].values   
    
    base_bert = AutoModelForSequenceClassification.from_pretrained(
        model_name,
        num_labels=len(label_encoder.classes_)
    )

    model = BertWithNumeric(base_bert.bert, num_labels=len(label_encoder.classes_))
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    
    model.to(device)

    train_dataset = BertLandUseDataset(train_texts, train_labels, tokenizer)    
    train_dataset1 = BertLandUseDataset(train_texts1, train_labels1, tokenizer)
    
    
    val_dataset = BertLandUseDataset(val_texts, val_labels, tokenizer)

    train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)   
    train_loader1 = DataLoader(train_dataset1, batch_size=16, shuffle=True)
     
    val_loader = DataLoader(val_dataset, batch_size=16)

    optimizer = AdamW(model.parameters(), lr=2e-5)
    criterion = torch.nn.CrossEntropyLoss()

    
    epochs = 5
    for epoch in range(epochs):

        print(f'Epoch {epoch+1} started')
        model.train()
        total_loss = 0
        for batch in train_loader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)

            optimizer.zero_grad()
            outputs = model(input_ids, attention_mask)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        print(f'Epoch {epoch+1}, Loss: {total_loss/len(train_loader)}')
        
    for epoch in range(epochs):

        print(f'Epoch {epoch+1} started')
        model.train()
        total_loss = 0
        for batch in train_loader1:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)

            optimizer.zero_grad()
            outputs = model(input_ids, attention_mask)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        print(f'Epoch {epoch+1}, Loss: {total_loss/len(train_loader1)}')
        
    save_full_model(model, tokenizer, label_encoder, text_columns, os_path, components_path)

else:

    
    predictor = LandUsePredictor(
        model_path=model_directory,
        components_path=model_components_directory
    )

    
path_to_csv_file = 'files/predictions_new.xlsx'

predictor.predict_data_from_client_csv_file(path_to_client_csv_file=excel_df_intermed_dir, sheet_name='данные', path_to_saved_file=path_to_csv_file)