Update app.py

app.py

@@ -1,4 +1,81 @@
 #허깅페이스에서 돌아갈 수 있도록 바꾸어 보았음
+import torch
+from transformers import BertTokenizerFast, BertForQuestionAnswering, Trainer, TrainingArguments
+from datasets import load_dataset
+from collections import defaultdict
+
+# 데이터 불러오기
+dataset_load = load_dataset('Multimodal-Fatima/OK-VQA_train')
+dataset = dataset_load['train'].select(range(300))
+
+# 불필요한 특성 선택
+selected_features = ['image', 'answers', 'question']
+selected_dataset = dataset.map(lambda ex: {feature: ex[feature] for feature in selected_features})
+
+# 소프트 인코딩
+answers_to_id = defaultdict(lambda: len(answers_to_id))
+selected_dataset = selected_dataset.map(lambda ex: {
+    'answers': [answers_to_id[ans] for ans in ex['answers']],
+    'question': ex['question'],
+    'image': ex['image']
+})
+
+id_to_answers = {v: k for k, v in answers_to_id.items()}
+id_to_labels = {k: ex['answers'] for k, ex in enumerate(selected_dataset)}
+
+selected_dataset = selected_dataset.map(lambda ex: {'answers': id_to_labels.get(ex['answers'][0]),
+                                                   'question': ex['question'],
+                                                   'image': ex['image']})
+
+flattened_features = []
+
+for ex in selected_dataset:
+    flattened_example = {
+        'answers': ex['answers'],
+        'question': ex['question'],
+        'image': ex['image'],
+    }
+    flattened_features.append(flattened_example)
+
+# 모델 가져오기
+from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer
+
+model_name = 'microsoft/git-base-vqav2'
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+
+# Trainer를 사용하여 모델 학습
+tokenizer = BertTokenizerFast.from_pretrained('bert-base-multilingual-cased')
+
+def preprocess_function(examples):
+    tokenized_inputs = tokenizer(examples['question'], truncation=True, padding=True)
+    return {
+        'input_ids': tokenized_inputs['input_ids'],
+        'attention_mask': tokenized_inputs['attention_mask'],
+        'pixel_values': [(4, 3, 244, 244)] * len(tokenized_inputs['input_ids']),
+        'pixel_mask': [1] * len(tokenized_inputs['input_ids']),
+        'labels': [[label] for label in examples['answers']]
+    }
+
+dataset = load_dataset("Multimodal-Fatima/OK-VQA_train")['train'].select(range(300))
+ok_vqa_dataset = dataset.map(preprocess_function, batched=True)
+ok_vqa_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'pixel_values', 'pixel_mask', 'labels'])
+
+training_args = TrainingArguments(
+    output_dir='./results',
+    num_train_epochs=20,
+    per_device_train_batch_size=4,
+    logging_steps=500,
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=ok_vqa_dataset
+)
+
+# 모델 학습
+trainer.train()
+
 
 import gradio as gr
 import torch