Jokh/disaster_class

import torch from transformers import AutoTokenizer, AutoModelForSequenceClassification

모델 및 토크나이저 로드

model_name = "jokh/disaster_class" labels = ["해당 없음", "1단계", "2단계", "3단계", "4단계", "5단계"]

모델과 토크나이저 불러오기

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=len(labels))

GPU 사용 가능 여부 확인 및 모델 로드

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") try: model.to(device) # 모델을 GPU 또는 CPU로 이동 model.eval() print("모델 로드 성공!") except RuntimeError as e: print(f"모델 로드 중 오류 발생: {e}") exit() # 오류 발생 시 프로그램 종료

예측 함수

def predict(text): encoding = tokenizer(text, padding="max_length", truncation=True, max_length=128, return_tensors="pt") encoding = encoding.to(device) # 입력 데이터를 모델과 동일한 장치로 이동 with torch.no_grad(): output = model(**encoding) prediction = torch.argmax(output.logits, dim=1).item() return labels[prediction]

사용자 입력 반복 실행

while True: text = input("신고 내용을 입력하세요 (종료하려면 'exit' 입력): ") if text.lower() == "exit": break result = predict(text) print(f"예측된 재난 등급: {result}")