Update train_model.py

train_model.py

@@ -15,62 +15,50 @@ if not hf_token:
 # Fazer login no Hugging Face
 subprocess.run(["huggingface-cli", "login", "--token", hf_token])
 
-# Clonar o repositório do dataset
-subprocess.run(["git", "clone", "https://github.com/DigantaD/NER.git"])
-data = pd.read_csv('NER/ner_dataset.csv', encoding='latin1')
-data = data.fillna(method='ffill')
+# Carregar os dados do dataset
+data = pd.read_csv('NER/ner_dataset.csv', encoding='latin1').fillna(method='ffill')
 
+# Preparar os dados
 unique_labels = data['Tag'].unique().tolist()
-
-class NERDataset:
-    def __init__(self, dataframe):
-        self.dataframe = dataframe
-        self.tokenizer = BertTokenizerFast.from_pretrained('bert-base-cased')
-        self.label2id = {label: i for i, label in enumerate(unique_labels)}
-        self.id2label = {i: label for label, i in self.label2id.items()}
-
-    def tokenize_and_align_labels(self, examples):
-        tokenized_inputs = self.tokenizer(examples['words'], truncation=True, is_split_into_words=True)
-        labels = []
-        for i, label in enumerate(examples['tags']):
-            word_ids = tokenized_inputs.word_ids(batch_index=i)
-            previous_word_idx = None
-            label_ids = []
-            for word_idx in word_ids:
-                if word_idx is None:
-                    label_ids.append(-100)
-                elif word_idx != previous_word_idx:
-                    label_ids.append(self.label2id[label[word_idx]])
-                else:
-                    label_ids.append(-100)
-                previous_word_idx = word_idx
-            labels.append(label_ids)
-        tokenized_inputs['labels'] = labels
-        return tokenized_inputs
-
-    def create_dataset(self):
-        grouped = self.dataframe.groupby('Sentence #').apply(lambda s: [(w, t) for w, t in zip(s['Word'].values.tolist(), s['Tag'].values.tolist())])
-        grouped = grouped.apply(pd.Series).reset_index()
-        grouped.columns = ['Sentence #', 'words_and_tags']
-
-        grouped['words'] = grouped['words_and_tags'].apply(lambda x: [w for w, t in x])
-        grouped['tags'] = grouped['words_and_tags'].apply(lambda x: [t for w, t in x])
-
-        dataset = Dataset.from_pandas(grouped[['words', 'tags']])
-        dataset = dataset.map(self.tokenize_and_align_labels, batched=True)
-        return dataset
-
-# Instanciar o dataset
-ner_dataset = NERDataset(data)
-dataset = ner_dataset.create_dataset()
+label2id = {label: i for i, label in enumerate(unique_labels)}
+
+def tokenize_and_align_labels(examples):
+    tokenized_inputs = tokenizer(examples['words'], truncation=True, is_split_into_words=True)
+    labels = []
+    for i, label in enumerate(examples['tags']):
+        word_ids = tokenized_inputs.word_ids(batch_index=i)
+        previous_word_idx = None
+        label_ids = []
+        for word_idx in word_ids:
+            if word_idx is None:
+                label_ids.append(-100)
+            elif word_idx != previous_word_idx:
+                label_ids.append(label2id[label[word_idx]])
+            else:
+                label_ids.append(-100)
+            previous_word_idx = word_idx
+        labels.append(label_ids)
+    tokenized_inputs['labels'] = labels
+    return tokenized_inputs
+
+grouped = data.groupby('Sentence #').apply(lambda s: [(w, t) for w, t in zip(s['Word'].values.tolist(), s['Tag'].values.tolist())])
+grouped = grouped.apply(pd.Series).reset_index()
+grouped.columns = ['Sentence #', 'words_and_tags']
+
+grouped['words'] = grouped['words_and_tags'].apply(lambda x: [w for w, t in x])
+grouped['tags'] = grouped['words_and_tags'].apply(lambda x: [t for w, t in x])
+
+dataset = Dataset.from_pandas(grouped[['words', 'tags']])
+tokenizer = BertTokenizerFast.from_pretrained('bert-base-cased')
+dataset = dataset.map(tokenize_and_align_labels, batched=True)
 
 # Dividir o dataset em treino e teste
 dataset = dataset.train_test_split(test_size=0.1)
 
-# Carregar o modelo pré-treinado
+# Carregar o modelo pré-treinado do Hugging Face
 model = BertForTokenClassification.from_pretrained('bert-base-cased', num_labels=len(unique_labels))
 
-# Treinamento
+# Definir argumentos de treinamento
 training_args = TrainingArguments(
     output_dir='./results',
     evaluation_strategy="epoch",
@@ -81,6 +69,7 @@ training_args = TrainingArguments(
     weight_decay=0.01,
 )
 
+# Inicializar o Trainer
 trainer = Trainer(
     model=model,
     args=training_args,
@@ -88,18 +77,7 @@ trainer = Trainer(
     eval_dataset=dataset['test'],
 )
 
+# Treinar o modelo
 trainer.train()
 
-# Salvar o modelo
-model.save_pretrained('./ner_model')
-ner_dataset.tokenizer.save_pretrained('./ner_model')
-
-# Verificar se o diretório do modelo foi criado corretamente
-model_dir = './ner_model'
-if os.path.exists(model_dir) and os.path.isdir(model_dir):
-    print(f"Diretório do modelo encontrado: {model_dir}")
-    print("Arquivos no diretório do modelo:")
-    for file_name in os.listdir(model_dir):
-        print(file_name)
-else:
-    print(f"Diretório do modelo não encontrado: {model_dir}")
+print("Treinamento do modelo concluído.")