mod

hft.py

@@ -102,9 +102,31 @@ def tokenize_function(examples):
         return_tensors="pt"
     )
     tokenized["labels"] = tokenized["input_ids"].clone()
-    tokenized["source_idx"] = examples["source_idx"]
+    tokenized["source_idx"] = torch.tensor(examples["source_idx"], dtype=torch.long)  # Upewnij się, że to tensor
     return tokenized
 
+def custom_collate_fn(batch):
+    input_ids = torch.stack([torch.tensor(b["input_ids"]) for b in batch])
+    attention_mask = torch.stack([torch.tensor(b["attention_mask"]) for b in batch])
+    labels = torch.stack([torch.tensor(b["labels"]) for b in batch])
+    source_idx = torch.tensor([b["source_idx"] for b in batch], dtype=torch.long)
+
+    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels, "source_idx": source_idx}
+
+def forward(self, input_ids=None, attention_mask=None, labels=None, source_idx=None, **kwargs):
+    outputs = super().forward(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        labels=labels,
+        **kwargs
+    )
+    
+    if source_idx is not None and source_idx.dim() == 1:
+        source_embeds = self.source_embedding(source_idx).unsqueeze(1)  # (batch_size, 1, hidden_size)
+        outputs.logits += source_embeds  # Upewnij się, że wymiary się zgadzają
+    
+    return outputs
+
 class CustomModel(AutoModelForCausalLM):
     def __init__(self, config):
         super().__init__(config)
@@ -171,6 +193,7 @@ trainer = CustomTrainer(
     model=model,
     args=training_args,
     train_dataset=tokenized_dataset,
+    data_collator=custom_collate_fn  # Dodaj customową funkcję collate
 )
 trainer.train()