mod

hft.py

@@ -145,28 +145,6 @@ class CustomTrainer(Trainer):
         loss = outputs.loss
         return (loss, outputs) if return_outputs else loss
 
-def generate_answer(question, model, tokenizer, source_mapper, max_length=200):
-    device = next(model.parameters()).device
-    inputs = tokenizer(question, return_tensors="pt", truncation=True, max_length=512)
-    inputs = {k: v.to(device) for k, v in inputs.items()}
-    
-    outputs = model.base_model.generate(
-        **inputs,
-        max_length=max_length,
-        num_return_sequences=1,
-        return_dict_in_generate=True,
-        output_scores=True,
-    )
-    
-    answer = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
-    
-    # Pobierz źródło z ostatniego tokena
-    last_token_id = outputs.sequences[0][-1].item()
-    source_idx = model.source_embedding.weight.shape[0] - 1
-    source = source_mapper.get_source(source_idx)
-    
-    return f"{answer}\n\nŹródło: {source if source else 'Opracowanie własne'}"
-
 # Inicjalizacja komponentów
 source_mapper = SourceMapper()
 model_name = "crumb/nano-mistral"
@@ -210,9 +188,28 @@ trainer = CustomTrainer(
 )
 trainer.train()
 
-# Przykładowe użycie
-model.eval()
+# Funkcja generująca odpowiedź
+def generate_answer(question, model, tokenizer, source_mapper, max_length=200):
+    inputs = tokenizer(question, return_tensors="pt", truncation=True, max_length=512)
+    
+    outputs = model.generate(
+        **inputs,
+        max_length=max_length,
+        num_return_sequences=1,
+        return_dict_in_generate=True,
+        output_scores=True,
+    )
+    
+    answer = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
+    
+    # Pobierz źródło z ostatniego tokena
+    last_token_id = outputs.sequences[0][-1].item()
+    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("Pytanie:", question)
-print("Odpowiedź:", answer)
+print(answer)