Create train.py

train.py

@@ -0,0 +1,29 @@
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
+from datasets import Dataset
+import json
+
+with open('train_data.json', 'r') as f:
+    data = json.load(f)
+
+texts = []
+labels = []
+for label, samples in data.items():
+    for text in samples:
+        texts.append(text)
+        labels.append(label)
+
+dataset = Dataset.from_dict({"text": texts, "label": labels})
+
+model_name = "distilbert-base-uncased"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=3)
+
+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=2)
+trainer = Trainer(model=model, args=training_args, train_dataset=tokenized_datasets)
+trainer.train()
+model.save_pretrained("./trained_model")