Update training/train_language.py

training/train_language.py

@@ -12,22 +12,27 @@ from language.intent import IntentClassifier
 from datasets import load_dataset
 
 # ================================
-# LOAD DATA FROM HUGGING FACE
+# CONFIG
 # ================================
-print("[INFO] Loading dataset from Hugging Face...")
-
-load_dataset("clinc_oos", "plus")
+ARTIFACTS_DIR = "artifacts"
+BATCH_SIZE = 16
+EPOCHS = 10
+LEARNING_RATE = 3e-4
+MAX_SEQ_LEN = 64
 
+os.makedirs(ARTIFACTS_DIR, exist_ok=True)
 
+# ================================
+# LOAD DATA FROM HUGGING FACE
+# ================================
+print("[INFO] Loading dataset from Hugging Face...")
 hf_dataset = load_dataset("clinc_oos", "plus")
 
 texts = hf_dataset["train"]["text"]
 labels = hf_dataset["train"]["intent"]
-
 intent_labels = sorted(list(set(labels)))
 
-print(f"[INFO] Loaded {len(texts)} samples from Hugging Face")
-
+print(f"[INFO] Loaded {len(texts)} samples with {len(intent_labels)} intents")
 
 # ================================
 # DATASET
@@ -43,9 +48,7 @@ class LanguageDataset(Dataset):
         return len(self.texts)
 
     def __getitem__(self, idx):
-        token_ids = self.tokenizer.encode(self.texts[idx])
-        # Truncate to max_seq_len
-        token_ids = token_ids[:self.max_seq_len]
+        token_ids = self.tokenizer.encode(self.texts[idx])[:self.max_seq_len]
         token_ids = torch.tensor(token_ids, dtype=torch.long)
         label = torch.tensor(self.labels[idx], dtype=torch.long)
         return token_ids, label
@@ -59,34 +62,11 @@ def collate_fn(batch, tokenizer):
     padded = []
     for t in token_ids:
         pad_len = max_len - len(t)
-        padded.append(
-            torch.cat([t, torch.full((pad_len,), tokenizer.vocab[tokenizer.PAD_TOKEN], dtype=torch.long)])
-        )
+        padded.append(torch.cat([t, torch.full((pad_len,), tokenizer.vocab[tokenizer.PAD_TOKEN], dtype=torch.long)]))
     return torch.stack(padded), torch.tensor(labels)
 
 # ================================
-# LOAD DATA
-# ================================
-def load_data(path):
-    texts, labels = [], []
-    intent_labels = set()
-    if not os.path.exists(path):
-        raise FileNotFoundError(f"Dataset file not found: {path}")
-
-    with open(path, "r", encoding="utf-8") as f:
-        for line in f:
-            line = line.strip()
-            if line:
-                text, intent = line.split("\t")
-                texts.append(text)
-                labels.append(intent)
-                intent_labels.add(intent)
-    return texts, labels, sorted(list(intent_labels))
-
-texts, labels, intent_labels = load_data("musombi/intent_datasets")
-
-# ================================
-# TOKENIZER
+# TOKENIZER AND DATALOADER
 # ================================
 tokenizer = SimpleTokenizer()
 tokenizer.build_vocab(texts)
@@ -95,8 +75,6 @@ tokenizer.freeze_vocab()
 dataset = LanguageDataset(texts, labels, tokenizer, intent_labels, max_seq_len=MAX_SEQ_LEN)
 loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=lambda batch: collate_fn(batch, tokenizer))
 
-print(f"[INFO] Loaded {len(dataset)} samples with {len(intent_labels)} intents")
-
 # ================================
 # MODEL
 # ================================
@@ -107,10 +85,7 @@ classifier = IntentClassifier(input_dim=encoder.projection.out_features, intent_
 embedder, encoder, classifier = embedder.to(DEVICE), encoder.to(DEVICE), classifier.to(DEVICE)
 
 criterion = nn.CrossEntropyLoss()
-optimizer = optim.Adam(
-    list(embedder.parameters()) + list(encoder.parameters()) + list(classifier.parameters()),
-    lr=LEARNING_RATE
-)
+optimizer = optim.Adam(list(embedder.parameters()) + list(encoder.parameters()) + list(classifier.parameters()), lr=LEARNING_RATE)
 
 # ================================
 # TRAINING LOOP
@@ -121,7 +96,6 @@ def train():
         total_loss = 0
         for token_ids, labels_batch in loader:
             token_ids, labels_batch = token_ids.to(DEVICE), labels_batch.to(DEVICE)
-
             embeddings = embedder(token_ids)
             attention_mask = (token_ids != tokenizer.vocab[tokenizer.PAD_TOKEN]).long()
             sentence_vec = encoder(embeddings, attention_mask=attention_mask)