Update app.py

app.py

@@ -1,34 +1,44 @@
 import gradio as gr
 import threading
+import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
 from datasets import load_dataset
 
-# 데이터셋 로딩
+# GPU가 아닌 CPU에서 실행하도록 설정
+device = torch.device("cpu")
+
+# IMDb 데이터셋 로딩
 dataset = load_dataset("imdb")
 
+# 데이터셋의 텍스트 컬럼 자동 감지
+text_column = dataset["train"].column_names[0]  # 기본적으로 "text"일 가능성이 높음
+
 # 모델과 토크나이저 로딩
 model_name = "distilbert-base-uncased"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
+model.to(device)  # 모델을 CPU로 이동
 
 # 데이터셋을 모델에 맞게 전처리
 def tokenize_function(examples):
-    return tokenizer(examples["text"], padding="max_length", truncation=True)
+    return tokenizer(examples[text_column], padding="max_length", truncation=True)
 
-tokenized_train_datasets = dataset["train"].map(tokenize_function, batched=True)
-tokenized_test_datasets = dataset["test"].map(tokenize_function, batched=True)
+tokenized_train_datasets = dataset["train"].map(tokenize_function, batched=True, batch_size=None, remove_columns=[text_column])
+tokenized_test_datasets = dataset["test"].map(tokenize_function, batched=True, batch_size=None, remove_columns=[text_column])
 
-# 훈련 설정 (빠르게 훈련하기 위해 에폭 수를 줄임)
+# 훈련 설정 (GPU 사용 안 함)
 training_args = TrainingArguments(
     output_dir="./results",           # 결과 저장 경로
     num_train_epochs=1,               # 훈련 에폭 수 1로 설정 (빠르게 테스트)
-    per_device_train_batch_size=16,   # 배치 크기 증가
-    per_device_eval_batch_size=16,    # 배치 크기 증가
+    per_device_train_batch_size=4,    # 배치 크기 줄이기 (CPU에서는 작은 값 추천)
+    per_device_eval_batch_size=4,     # 배치 크기 줄이기
     evaluation_strategy="epoch",      # 에폭마다 검증
+    save_strategy="epoch",
     logging_dir="./logs",             # 로그 저장 경로
     logging_steps=100,                # 100 스텝마다 로그 출력
-    report_to="tensorboard",          # 텐서보드로 로그 보고
+    report_to="none",                 # 허깅페이스 업로드 시 로깅 비활성화
     load_best_model_at_end=True,      # 최상의 모델로 종료
+    no_cuda=True                      # ❌ GPU 사용하지 않도록 설정
 )
 
 # 훈련 함수
@@ -46,25 +56,23 @@ def start_training():
     train_thread = threading.Thread(target=train_model)
     train_thread.start()
 
-# 그라디언트 기반 훈련된 모델을 UI에 연결
+# 텍스트 분류 함수 (CPU에서 실행)
 def classify_text(text):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    outputs = model(**inputs)
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
+    with torch.no_grad():  # 불필요한 연산 방지
+        outputs = model(**inputs)
     logits = outputs.logits
     predicted_class = logits.argmax(-1).item()
-    return predicted_class
+    return str(predicted_class)  # Gradio에서 문자열 반환이 더 안정적
 
 # Gradio 인터페이스 설정
 demo = gr.Interface(fn=classify_text, inputs="text", outputs="text")
 
 # 훈련 시작과 Gradio UI 실행
 def launch_app():
-    # 훈련을 시작
-    start_training()
-
-    # Gradio 인터페이스 실행
-    demo.launch()
+    start_training()  # 훈련 시작
+    demo.launch()     # Gradio UI 실행
 
-# 허깅페이스 Spaces에 업로드 할 때는 이 부분을 실행하도록 설정
+# 허깅페이스 Spaces에 업로드할 때 실행
 if __name__ == "__main__":
     launch_app()