dif

app.py

@@ -1,8 +1,50 @@
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
+from datasets import load_dataset
 
-def greet(name):
-    return "Hello " + name + "!!"
+# 데이터셋 로딩
+dataset = load_dataset("imdb")
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-print("doneasdf")
+# 모델과 토크나이저 로딩
+model_name = "distilbert-base-uncased"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+
+# 데이터셋을 모델에 맞게 전처리
+def tokenize_function(examples):
+    return tokenizer(examples["text"], padding="max_length", truncation=True)
+
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
+
+# 훈련 설정
+training_args = TrainingArguments(
+    output_dir="./results",           # 결과 저장 경로
+    num_train_epochs=3,               # 훈련 에폭 수
+    per_device_train_batch_size=8,    # 배치 크기
+    per_device_eval_batch_size=8,     # 검증 배치 크기
+    evaluation_strategy="epoch",      # 에폭마다 검증
+    logging_dir="./logs",             # 로그 저장 경로
+)
+
+trainer = Trainer(
+    model=model,                       # 훈련할 모델
+    args=training_args,                # 훈련 인자
+    train_dataset=tokenized_datasets["train"],  # 훈련 데이터셋
+    eval_dataset=tokenized_datasets["test"],    # 평가 데이터셋
+)
+
+# 훈련 시작
+trainer.train()
+
+# 그라디오 인터페이스로 훈련된 모델을 UI에 연결
+def classify_text(text):
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class = logits.argmax(-1).item()
+    return predicted_class
+
+demo = gr.Interface(fn=classify_text, inputs="text", outputs="text")
+
+# Gradio 인터페이스 실행 (훈련 후)
 demo.launch()