[빅프로젝트] 3주차의 시행착오 - 영문 텍스트화 및 text classification 모델 구축

https://www.aihub.or.kr/aihubdata/data/view.do?currMenu=115&topMenu=100&aihubDataSe=data&dataSetSn=71768

AI-Hub

 

원래 우리 프로젝트에 사용하려고 했던 데이터셋이다.

 

그러나 1차 개방에 포함되지 않아서 문의해보니, 7월달에 추가되는 건 어렵겠다고 직접 전화가 오셨다.

빅프 발표가 7월 31일이라 결국 우리는 대체를 찾아야 했다.

그나마 대응을 할 수 있도록 해주셔서 감사할 따름이다.

 

따라서 우리 프로젝트는 일단 영어로 된 데이터셋으로 우선 구축하고, 해당 텍스트 셋을 이용해서 시연을 하기로 코치님께 허락받았다.

https://www.kaggle.com/datasets/louisteitelbaum/911-recordings

 

Recording 파일에 텍스트 데이터는 없지만, 내가 또 누구인가. 무려 영어권에서 살다 온 사람(사실 모델로 돌려서 상관 없지만)

 

Whisper-large-v2 모델로 돌려 해당 recording 파일을 텍스트로 바꾸는 작업을 수행했다.

 

이때 어시스턴트 모델을 쓰면 속도가 2배 빨라진다는 깃허브의 코드를 보고, 코드를 짜서 전부 txt 파일로 바꿨다.

import os
import torch
from transformers import AutoModelForCausalLM, AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
from tqdm import tqdm

# GPU 사용 가능 여부 확인 및 장치 설정
device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

# 메인 모델 로드
model_id = "openai/whisper-large-v2"
model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)

# 보조 모델 로드
assistant_model_id = "distil-whisper/distil-large-v2"
assistant_model = AutoModelForCausalLM.from_pretrained(
    assistant_model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
assistant_model.to(device)

# 프로세서 로드
processor = AutoProcessor.from_pretrained(model_id)

# 파이프라인 인스턴스화
pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    max_new_tokens=128,
    generate_kwargs={"assistant_model": assistant_model},
    torch_dtype=torch_dtype,
    device=device,
)

# Kaggle input 폴더 경로 설정
input_folder = "/kaggle/input/911-recordings-wav/cleaned_audio"  # 실제 데이터셋 이름으로 변경해주세요
output_folder = "/kaggle/working/transcriptions"  # 출력 폴더 경로

# 출력 폴더가 없으면 생성
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# wav 파일 목록 가져오기
wav_files = [f for f in os.listdir(input_folder) if f.endswith('.wav')]

# 변환 및 개별 파일로 저장
for wav_file in tqdm(wav_files, desc="Processing audio files"):
    input_path = os.path.join(input_folder, wav_file)
    output_path = os.path.join(output_folder, f"{os.path.splitext(wav_file)[0]}.txt")
    
    result = pipe(input_path)
    
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(result['text'])

print(f"Transcription completed. Results saved in {output_folder}")

 

 

확실히 영어 recording이라 그런지, 위스퍼 모델이 진짜 엄청난 텍스트화를 보여주는 기염을 토했다.

 

그 다음에는 긴급도 분류.

 

해당 데이터셋에 메타데이터가 있긴 하지만, 어떻게 이걸 정해야 할지 난감하긴 하다.

그러나 일단 사상자와 description에 따라서 상/중/하로 분류하는 메커니즘을 만드는 것 까지가 월요일에 한 일.

 

영문 sequence classification을 위해서 일단 여러가지 bert 모델도 써볼 생각이다.

 

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification, Trainer, TrainingArguments
from datasets import Dataset

# 데이터 준비 (CSV 파일을 읽어와서 pandas DataFrame으로 변환)
df = pd.read_csv('path_to_your_dataset.csv')

# 필요한 컬럼만 선택
df = df[['text', 'label']]  # text와 label 컬럼이 있다고 가정

# Train, Validation, Test 데이터셋 분리
train_df, temp_df = train_test_split(df, test_size=0.3, stratify=df['label'])
val_df, test_df = train_test_split(temp_df, test_size=0.5, stratify=temp_df['label'])

# pandas DataFrame을 datasets Dataset으로 변환
train_dataset = Dataset.from_pandas(train_df)
val_dataset = Dataset.from_pandas(val_df)
test_dataset = Dataset.from_pandas(test_df)

# 토크나이저와 모델 불러오기
model_name = "distilbert-base-uncased"
tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
model = DistilBertForSequenceClassification.from_pretrained(model_name, num_labels=3)

# 데이터셋 전처리
def preprocess_function(examples):
    return tokenizer(examples['text'], truncation=True, padding=True)

train_dataset = train_dataset.map(preprocess_function, batched=True)
val_dataset = val_dataset.map(preprocess_function, batched=True)
test_dataset = test_dataset.map(preprocess_function, batched=True)

# Trainer를 사용하여 모델 학습
training_args = TrainingArguments(
    output_dir='./results',
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=3,
    weight_decay=0.01,
    logging_dir='./logs',
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    tokenizer=tokenizer,
    compute_metrics=lambda p: {'accuracy': accuracy_score(p.label_ids, p.predictions.argmax(-1))}
)

trainer.train()

# 모델 평가
results = trainer.evaluate(eval_dataset=test_dataset)
print(f"Test Accuracy: {results['eval_accuracy']}")

# 모델 저장
model.save_pretrained('./emergency_classification_model')
tokenizer.save_pretrained('./emergency_classification_model')

 

이 프로젝트에 대한 백엔드 플로우도 다 같이 짜야 하니, 오늘은 이만 줄이겠다.