From fef5717c2b5d1f622a1883708e8083c2cf0eeda1 Mon Sep 17 00:00:00 2001
From: "l.gabrysiak" <l.gabrysiak@osadkowski.pl>
Date: Tue, 25 Feb 2025 16:07:46 +0100
Subject: [PATCH] zmiany

---
 hft.py | 35 ++++++++++++++++++++---------------
 1 file changed, 20 insertions(+), 15 deletions(-)

diff --git a/hft.py b/hft.py
index 2ab3285..d69af70 100644
--- a/hft.py
+++ b/hft.py
@@ -109,22 +109,26 @@ 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])
-    
-    # Dodajemy domyślne source_idx, jeśli nie istnieje
     source_idx = torch.tensor([b.get("source_idx", -1) for b in batch], dtype=torch.long)
-    print("source_idx shape:", source_idx.shape)  # Debugowanie
-    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels, "source_idx": source_idx}
+
+    return {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "labels": labels,
+        "source_idx": source_idx
+    }
 
 class CustomModel(AutoModelForCausalLM):
     def __init__(self, config):
         super().__init__(config)
         self.source_embedding = nn.Embedding(
-            num_embeddings=1000,  # Maksymalna liczba unikalnych źródeł
+            num_embeddings=1000,
             embedding_dim=config.hidden_size,
             padding_idx=-1
         )
         
-    def forward(self, input_ids=None, attention_mask=None, labels=None, source_idx=None, **kwargs):
+    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
+        source_idx = kwargs.pop('source_idx', None)  # Pobierz i usuń source_idx z kwargs
         outputs = super().forward(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -133,24 +137,25 @@ class CustomModel(AutoModelForCausalLM):
         )
         
         if source_idx is not None:
-            # Tutaj dodaj logikę obsługi source_idx
-            pass
-        
+            source_idx = source_idx.to(outputs.logits.device)  # Ensure same device
+            source_embeds = self.source_embedding(source_idx).unsqueeze(1)
+            outputs.logits += source_embeds
+            
         return outputs
 
-
 class CustomTrainer(Trainer):
-    def compute_loss(self, model, inputs, return_outputs=False):
+    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
         labels = inputs.pop("labels")
-        source_idx = inputs.pop("source_idx", None)
-        outputs = model(**inputs, labels=labels, source_idx=source_idx if source_idx is not None else None)
+        source_idx = inputs.pop("source_idx")
+        outputs = model(**inputs, labels=labels, source_idx=source_idx)
         return (outputs.loss, outputs) if return_outputs else outputs.loss
 
 # Inicjalizacja komponentów
 source_mapper = SourceMapper()
 model_name = "crumb/nano-mistral" #"google/gemma-2-2b"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-tokenizer.pad_token = tokenizer.eos_token
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
 
 # Przygotowanie danych
 catalog_path = "file_catalog.json"
@@ -161,7 +166,7 @@ tokenized_dataset = dataset.map(tokenize_function, batched=True, batch_size=8)
 # Inicjalizacja modelu
 config = AutoModelForCausalLM.from_pretrained(model_name).config
 model = CustomModel.from_pretrained(model_name, config=config)
-model.to("cpu")
+model.to("cuda" if torch.cuda.is_available() else "cpu")
 
 # Konfiguracja treningu
 training_args = TrainingArguments(