This commit is contained in:
l.gabrysiak 2025-02-26 00:08:31 +01:00
parent 4be750503d
commit 4957a2898b
1 changed files with 13 additions and 17 deletions

30
hft.py
View File

@ -103,21 +103,19 @@ def prepare_dataset(directory, catalog_path, source_mapper):
current_section = "" current_section = ""
current_chapter = "" current_chapter = ""
# Wykrywanie struktury dokumentu
structure_matches = re.finditer( structure_matches = re.finditer(
r'(DZIAŁ [A-ZĄĆĘŁŃÓŚŹŻ]+)\n+(.*?)\n(?=Art\.|Rozdział|DZIAŁ|$)' r'(DZIAŁ [A-ZĄĆĘŁŃÓŚŹŻ]+)\n+(.*?)\n(?=Art\.|Rozdział|DZIAŁ|$)'
r'|(Rozdział [A-ZĄĆĘŁŃÓŚŹŻ]+)\n+(.*?)\n(?=Art\.|DZIAŁ|$)', r'|(Rozdział [A-ZĄĆĘŁŃÓŚŹŻ]+)\n+(.*?)\n(?=Art\.|DZIAŁ|$)',
text text
) )
for match in structure_matches: for match in structure_matches:
if match.group(1): # DZIAŁ if match.group(1):
current_section = f"{match.group(1)} - {match.group(2).strip()}" current_section = f"{match.group(1)} - {match.group(2).strip()}"
current_chapter = "" current_chapter = ""
else: # Rozdział else:
current_chapter = f"{match.group(3)} - {match.group(4).strip()}" current_chapter = f"{match.group(3)} - {match.group(4).strip()}"
if doc_type != "Opracowanie własne": if doc_type != "Opracowanie własne":
# Ulepszony regex dla artykułów
articles = re.split( articles = re.split(
r'(?i)(Art[\.\s]+\d+[a-z]*(?:[\\.-]\d+)*)\.?\s*', r'(?i)(Art[\.\s]+\d+[a-z]*(?:[\\.-]\d+)*)\.?\s*',
text text
@ -133,7 +131,6 @@ def prepare_dataset(directory, catalog_path, source_mapper):
if len(article_content) < 50: if len(article_content) < 50:
continue continue
# Formatowanie cytowania
citation_block = ( citation_block = (
f"{CITATION_START}\n" f"{CITATION_START}\n"
f"Dokument: {doc_type}\n" f"Dokument: {doc_type}\n"
@ -177,15 +174,15 @@ def prepare_dataset(directory, catalog_path, source_mapper):
return data return data
class CustomModel(nn.Module): class CustomModel(nn.Module):
def __init__(self, model_name, config): def __init__(self, model_name, tokenizer):
super().__init__() super().__init__()
config = AutoModelForCausalLM.from_pretrained(model_name).config
self.base_model = AutoModelForCausalLM.from_pretrained(model_name, config=config) self.base_model = AutoModelForCausalLM.from_pretrained(model_name, config=config)
self.source_embedding = nn.Embedding(10000, config.hidden_size, padding_idx=-1) self.source_embedding = nn.Embedding(10000, config.hidden_size, padding_idx=-1)
# Dodatkowa inicjalizacja tokenizera # Dodaj specjalne tokeny i zaktualizuj embeddings
self.tokenizer = AutoTokenizer.from_pretrained(model_name) tokenizer.add_special_tokens({'additional_special_tokens': [CITATION_START, CITATION_END]})
self.tokenizer.add_special_tokens({'additional_special_tokens': [CITATION_START, CITATION_END]}) self.base_model.resize_token_embeddings(len(tokenizer))
self.base_model.resize_token_embeddings(len(self.tokenizer))
for param in self.base_model.parameters(): for param in self.base_model.parameters():
param.requires_grad = False param.requires_grad = False
@ -218,7 +215,7 @@ class CustomDataCollator(DataCollatorForLanguageModeling):
batch = super().torch_call(examples) batch = super().torch_call(examples)
if "source_idx" in examples[0]: if "source_idx" in examples[0]:
source_idx = torch.stack([torch.tensor(ex["source_idx"]) for ex in examples]) source_idx = torch.tensor([ex["source_idx"] for ex in examples])
batch["source_idx"] = source_idx batch["source_idx"] = source_idx
return batch return batch
@ -226,12 +223,11 @@ class CustomDataCollator(DataCollatorForLanguageModeling):
def main(): def main():
source_mapper = SourceMapper() source_mapper = SourceMapper()
model_name = "crumb/nano-mistral" model_name = "crumb/nano-mistral"
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Dodaj specjalne tokeny do tokenizera # Inicjalizacja tokenizera
tokenizer.add_special_tokens({'additional_special_tokens': [CITATION_START, CITATION_END]}) tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token tokenizer.pad_token = tokenizer.eos_token
# Przygotowanie danych # Przygotowanie danych
catalog_path = "catalog.json" catalog_path = "catalog.json"
data = prepare_dataset("docs", catalog_path, source_mapper) data = prepare_dataset("docs", catalog_path, source_mapper)
@ -259,8 +255,8 @@ def main():
tokenized_dataset = dataset.map(tokenize_function, batched=True, batch_size=16) tokenized_dataset = dataset.map(tokenize_function, batched=True, batch_size=16)
model = CustomModel(model_name, AutoModelForCausalLM.from_pretrained(model_name).config) # Inicjalizacja modelu z tokenizerem
model.source_mapper = source_mapper model = CustomModel(model_name, tokenizer)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device) model.to(device)