Create train.py

train.py

@@ -0,0 +1,87 @@
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from datasets import load_dataset
+from torch.utils.data import Dataset, DataLoader
+import os
+
+# 加载数据集
+datasets = [
+    "Johnson8187/Chinese_Multi-Emotion_Dialogue_Dataset",
+    "clapAI/MultiLingualSentiment",
+    "shareAI/ShareGPT-Chinese-English-90k",
+    "wikimedia/wikipedia",
+    "google/code_x_glue_tt_text_to_text",
+    "silk-road/ChatHaruhi-54K-Role-Playing-Dialogue",
+    "yentinglin/TaiwanChat",
+    "liswei/rm-static-zhTW",
+    "yys/OpenOrca-Chinese",
+    "Fumika/Wikinews-multilingual",
+    "aqweteddy/Taiwan-Curlture-MCQ",
+    "Nexdata/Chinese_Mandarin_Multi-emotional_Synthesis_Corpus",
+    "Nexdata/Chinese_Mandarin_Entertainment_anchor_Style_Multi-emotional_Synthesis_Corpus",
+    "voices365/102_Hours_High_Quality_Chinese_Audio_Dataset_For_Speech_Synthesis_Female_Samples",
+    "voices365/Chinese_Female_001VoiceArtist_40Hours_High_Quality_Voice_Dataset",
+    "Nexdata/Mandarin_Spontaneous_Speech_Data",
+    "speechbrain/common_language",
+    "hello2mao/Chinese_Audio_Resource"
+]
+
+# 加载模型和tokenizer
+model = AutoModelForSequenceClassification.from_pretrained("zeroMN/zeroSG")
+tokenizer = AutoTokenizer.from_pretrained("zeroMN/zeroSG")
+
+# 创建数据加载器
+class MyDataset(Dataset):
+    def __init__(self, datasets):
+        self.datasets = datasets
+        self.data = []
+        for dataset in datasets:
+            data = load_dataset(dataset)
+            self.data.extend(data["train"])
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        text = self.data[idx]["text"]
+        inputs = tokenizer.encode_plus(
+            text,
+            add_special_tokens=True,
+            max_length=512,
+            return_attention_mask=True,
+            return_tensors='pt'
+        )
+        return {
+            'input_ids': inputs['input_ids'].flatten(),
+            'attention_mask': inputs['attention_mask'].flatten(),
+            'labels': torch.tensor(0)  # placeholder for labels
+        }
+
+dataset = MyDataset(datasets)
+data_loader = DataLoader(dataset, batch_size=32, shuffle=True)
+
+# 训练模型
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model.to(device)
+criterion = torch.nn.CrossEntropyLoss()
+optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
+
+for epoch in range(5):
+    model.train()
+    total_loss = 0
+    for batch in data_loader:
+        input_ids = batch['input_ids'].to(device)
+        attention_mask = batch['attention_mask'].to(device)
+        labels = batch['labels'].to(device)
+
+        optimizer.zero_grad()
+
+        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
+        loss = criterion(outputs, labels)
+
+        loss.backward()
+        optimizer.step()
+
+        total_loss += loss.item()
+
+    print(f'Epoch {epoch+1}, Loss: {total_loss / len(data_loader)}')