diff --git a/hft.py b/hft.py
index 8196610..0aedb93 100644
--- a/hft.py
+++ b/hft.py
@@ -204,9 +204,14 @@ def generate_answer(question, model, tokenizer, source_mapper, max_length=200):
     
     # Pobierz źródło z ostatniego tokena
     last_token_id = outputs.sequences[0][-1].item()
-    source_idx = model.source_embeddi
+    source_idx = model.source_embedding.weight.shape[0] - 1  # Tymczasowe rozwiązanie
+    source = source_mapper.get_source(source_idx)
+    
+    return f"{answer}\n\nŹródło: {source if source else 'Opracowanie własne'}"
 
 # Przykład użycia
 question = "Ile dni urlopu przysługuje pracownikowi?"
 answer = generate_answer(question, model, tokenizer, source_mapper)
-print(answer)
\ No newline at end of file
+print("Pytanie:", question)
+print("Odpowiedź:", answer)
+