diff --git a/hft.py b/hft.py index 18e53cb..c912e43 100644 --- a/hft.py +++ b/hft.py @@ -9,20 +9,36 @@ import pytesseract import docx2txt import PyPDF2 import json - +from collections import defaultdict from huggingface_hub import login -login(f"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") +login(token="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") os.environ["TOKENIZERS_PARALLELISM"] = "false" +# Nowa klasa do zarządzania źródłami +class SourceMapper: + def __init__(self): + self.source_to_idx = defaultdict(lambda: len(self.source_to_idx)) + self.idx_to_source = {} + + def add_source(self, source): + if source and source not in self.source_to_idx: + idx = self.source_to_idx[source] + self.idx_to_source[idx] = source + + def get_idx(self, source): + return self.source_to_idx[source] if source else -1 + + def get_source(self, idx): + return self.idx_to_source.get(idx, "Unknown") + def load_file_catalog(catalog_path): with open(catalog_path, 'r', encoding='utf-8') as file: return json.load(file) def identify_legal_document(filename, file_catalog): - return file_catalog.get(filename, f"") + return file_catalog.get(filename, "Opracowanie własne") -# Funkcja do ekstrakcji tekstu z różnych typów plików def extract_text_from_file(file_path): _, ext = os.path.splitext(file_path) ext = ext.lower() @@ -44,117 +60,142 @@ def extract_text_from_file(file_path): else: return "" -# Przygotowanie danych -def prepare_dataset(directory, catalog_path): +def prepare_dataset(directory, catalog_path, source_mapper): file_catalog = load_file_catalog(catalog_path) data = [] + for root, _, files in os.walk(directory): for file in files: file_path = os.path.join(root, file) text = extract_text_from_file(file_path) - if text: - # Sprawdzenie, czy plik znajduje się w katalogu - doc_type = identify_legal_document(file, file_catalog) - if doc_type != "Opracowanie własne": - # Przetwarzanie dla aktów prawnych - articles = re.split(r'(Art\.\s+\d+\.)', text)[1:] - for i in range(0, len(articles), 2): - if i + 1 < len(articles): - article_number = articles[i].strip() - article_content = articles[i + 1].strip() - data.append({ - "text": f"{article_number} {article_content}", - "source": f"{doc_type}, {article_number}" - }) - else: - # Przetwarzanie dla zwykłych dokumentów - chunks = [text[i:i + 512] for i in range(0, len(text), 512)] - for chunk in chunks: - data.append({ - "text": chunk, - "source": f"" - }) + if not text: + continue + + doc_type = identify_legal_document(file, file_catalog) + if doc_type != "Opracowanie własne": + articles = re.split(r'(Art\.\s+\d+[\.\s])', text) + for i in range(1, len(articles), 2): + article_number = articles[i].strip() + article_content = articles[i+1].strip() if i+1 < len(articles) else "" + source = f"{doc_type}, {article_number}" + source_mapper.add_source(source) + + data.append({ + "text": f"{article_number} {article_content}", + "source_idx": source_mapper.get_idx(source) + }) + else: + chunks = [text[i:i+512] for i in range(0, len(text), 512)] + for chunk in chunks: + data.append({ + "text": chunk, + "source_idx": -1 # Brak źródła + }) return data - -# Tokenizacja danych def tokenize_function(examples): - inputs = tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512) - inputs["labels"] = inputs["input_ids"].copy() - inputs["source"] = examples["source"] - return inputs + tokenized = tokenizer( + examples["text"], + truncation=True, + padding="max_length", + max_length=512, + return_tensors="pt" + ) + tokenized["labels"] = tokenized["input_ids"].clone() + tokenized["source_idx"] = examples["source_idx"] + return tokenized -# Dostosowany model class CustomModel(AutoModelForCausalLM): def __init__(self, config): super().__init__(config) - self.source_embedding = nn.Embedding(1000, config.hidden_size) # Zakładamy maksymalnie 1000 różnych źródeł - - def forward(self, input_ids, attention_mask=None, labels=None, source=None): - outputs = super().forward(input_ids, attention_mask=attention_mask, labels=labels) - if source is not None: - source_embeds = self.source_embedding(source) - outputs.logits += source_embeds.unsqueeze(1) + self.source_embedding = nn.Embedding( + num_embeddings=1000, # Maksymalna liczba unikalnych źródeł + embedding_dim=config.hidden_size, + padding_idx=-1 + ) + + 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: + # Dodajemy embedding źródła do hidden states + source_embeds = self.source_embedding(source_idx).unsqueeze(1) + outputs.logits += source_embeds + return outputs -# Dostosowany Trainer 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 = inputs.pop("source", None) # Użyj None jako wartości domyślnej - outputs = model(**inputs, labels=labels) - loss = outputs.loss - return (loss, outputs) if return_outputs else loss + 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 - -# Przygotowanie modelu i tokenizera +# Inicjalizacja komponentów +source_mapper = SourceMapper() model_name = "google/gemma-2-2b" tokenizer = AutoTokenizer.from_pretrained(model_name) -model = CustomModel.from_pretrained(model_name) +tokenizer.pad_token = tokenizer.eos_token -# Przygotowanie datasetu +# Przygotowanie danych catalog_path = "file_catalog.json" -data = prepare_dataset("files", catalog_path) +data = prepare_dataset("files", catalog_path, source_mapper) dataset = Dataset.from_list(data) -tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset.column_names) +tokenized_dataset = dataset.map(tokenize_function, batched=True, batch_size=32) + +# Inicjalizacja modelu +config = AutoModelForCausalLM.from_pretrained(model_name).config +model = CustomModel.from_pretrained(model_name, config=config) # Konfiguracja treningu training_args = TrainingArguments( output_dir="./results", num_train_epochs=3, - per_device_train_batch_size=4, - save_steps=10_000, - save_total_limit=2, + per_device_train_batch_size=2, + gradient_accumulation_steps=4, + learning_rate=2e-5, + fp16=True, + logging_steps=100, + save_strategy="steps", + save_steps=1000, + report_to="none" ) -# Inicjalizacja Trainera +# Trening trainer = CustomTrainer( model=model, args=training_args, train_dataset=tokenized_dataset, ) - -# Trening modelu trainer.train() -# Zapisanie modelu -trainer.save_model("./gemma2_finetuned") - -# Funkcja do generowania odpowiedzi z cytowaniem -def generate_answer(question, model, tokenizer, dataset): - inputs = tokenizer(question, return_tensors="pt") - outputs = model.generate(**inputs, max_length=200, num_return_sequences=1, output_scores=True, return_dict_in_generate=True) +# 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) - # Znajdź najbardziej prawdopodobne źródło - source_probs = outputs.scores[-1][:, model.source_embedding.weight.shape[0]:] - most_likely_source_idx = torch.argmax(source_probs).item() - most_likely_source = dataset[most_likely_source_idx % len(dataset)]['source'] + # 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: {most_likely_source}" + 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, dataset) +answer = generate_answer(question, model, tokenizer, source_mapper) print(answer) \ No newline at end of file