import os import torch import torch.nn as nn from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer from datasets import Dataset import re import json import PyPDF2 import docx2txt import pytesseract from PIL import Image from collections import defaultdict from huggingface_hub import login # Konfiguracja os.environ['TORCH_USE_CUDA_DSA'] = '1' os.environ["TOKENIZERS_PARALLELISM"] = "false" login(token="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") 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): try: with open(catalog_path, 'r', encoding='utf-8') as file: return json.load(file) except Exception as e: print(f"Błąd wczytywania katalogu plików: {str(e)}") return {} def identify_legal_document(filename, file_catalog): base_name = os.path.splitext(filename)[0].lower() return file_catalog.get(base_name, "Opracowanie własne") def extract_text_from_file(file_path): try: _, ext = os.path.splitext(file_path) ext = ext.lower() if ext in ['.txt', '.md']: with open(file_path, 'r', encoding='utf-8') as file: return file.read() elif ext == '.pdf': text = "" try: with open(file_path, 'rb') as file: reader = PyPDF2.PdfReader(file) for page in reader.pages: text += page.extract_text() or "" except Exception as e: print(f"Błąd PDF: {str(e)}") return text elif ext in ['.doc', '.docx']: return docx2txt.process(file_path) elif ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']: return pytesseract.image_to_string(Image.open(file_path)) else: print(f"Nieobsługiwany format pliku: {ext}") return "" except Exception as e: print(f"Błąd ekstrakcji tekstu: {str(e)}") return "" def prepare_dataset(directory, catalog_path, source_mapper): file_catalog = load_file_catalog(catalog_path) data = [] print(f"\n{'='*50}\nDIAGNOSTYKA DANYCH\n{'='*50}") for root, _, files in os.walk(directory): for file in files: file_path = os.path.join(root, file) print(f"\nPrzetwarzanie pliku: {file_path}") try: text = extract_text_from_file(file_path) if not text.strip(): print("Pominięto - brak tekstu") continue print(f"Długość tekstu: {len(text)} znaków") doc_type = identify_legal_document(file, file_catalog) print(f"Rozpoznany typ dokumentu: {doc_type}") if doc_type != "Opracowanie własne": articles = re.split(r'(?i)(Art[\.\s]+\d+[\.\s]?)', text) articles = [a.strip() for a in articles if a.strip()] print(f"Znaleziono {len(articles)} fragmentów") # Generowanie większej liczby przykładów for i in range(0, len(articles)-1, 2): for chunk_size in [256, 512, 1024]: # Różne rozmiary chunków article_number = articles[i] article_content = articles[i+1] chunks = [article_content[j:j+chunk_size] for j in range(0, len(article_content), chunk_size//2)] chunks = [c for c in chunks if len(c) > 100] for chunk in chunks: source = f"{doc_type}, {article_number}" source_mapper.add_source(source) data.append({ "text": f"{article_number} {chunk}", "source_idx": source_mapper.get_idx(source) }) else: clean_text = re.sub(r'\s+', ' ', text).strip() for chunk_size in [256, 512, 768]: # Trzy różne rozmiary chunks = [clean_text[i:i+chunk_size] for i in range(0, len(clean_text), chunk_size//2)] chunks = [c for c in chunks if c.strip()] for chunk in chunks: data.append({ "text": chunk, "source_idx": -1 }) print(f"Dodano {len(chunks)*3} chunków") except Exception as e: print(f"Błąd podczas przetwarzania pliku: {str(e)}") continue print(f"\nPodsumowanie przygotowania danych:") print(f"Łączna liczba przykładów: {len(data)}") if data: print("Przykładowy wpis:") print(json.dumps(data[0], indent=2, ensure_ascii=False)) else: print("BRAK DANYCH - sprawdź diagnostykę powyżej") return data class CustomModel(nn.Module): def __init__(self, model_name, config): super().__init__() self.base_model = AutoModelForCausalLM.from_pretrained(model_name, config=config) self.source_embedding = nn.Embedding(10000, config.hidden_size, padding_idx=-1) # Fine-tuning części modelu for param in self.base_model.parameters(): param.requires_grad = False for param in self.base_model.get_output_embeddings().parameters(): param.requires_grad = True for param in self.base_model.get_input_embeddings().parameters(): param.requires_grad = True def forward(self, input_ids=None, attention_mask=None, labels=None, source_idx=None, **kwargs): if source_idx is not None: valid_indices = torch.clamp(source_idx, 0, self.source_embedding.num_embeddings-1) source_embeds = self.source_embedding(valid_indices).unsqueeze(1) inputs_embeds = self.base_model.get_input_embeddings()(input_ids) + source_embeds return self.base_model(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels, **kwargs) return self.base_model(input_ids=input_ids, attention_mask=attention_mask, labels=labels, **kwargs) def generate(self, *args, **kwargs): return self.base_model.generate(*args, **kwargs) class CustomTrainer(Trainer): def __init__(self, *args, **kwargs): self.tokenizer = kwargs.pop('tokenizer', None) super().__init__(*args, **kwargs) 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) return (outputs.loss, outputs) if return_outputs else outputs.loss def evaluate(self): questions = [ "Jakie są prawa pracownika według art. 1?", "Kto jest pracownikiem według art. 2?", "Jakie są obowiązki pracodawcy według art. 3?" ] print("\n" + "="*50 + "\nEWALUACJA\n" + "="*50) for q in questions: result = self.generate_answer(q) print(f"\nPYTANIE: {q}") print(f"ODPOWIEDŹ: {result['answer'][:500]}") print(f"ŹRÓDŁA: {', '.join(result['sources'])}") print("-"*80) return {"loss": 0.0} def generate_answer(self, question): inputs = self.tokenizer( f"[PYTANIE] {question} [KONTEKST]", return_tensors="pt", truncation=True, max_length=512 ).to(self.model.base_model.device) with torch.no_grad(): outputs = self.model.generate( **inputs, max_new_tokens=200, temperature=0.5, top_p=0.9, repetition_penalty=2.0, num_beams=3, no_repeat_ngram_size=3 ) answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True) answer = answer.split("[KONTEKST]")[-1].strip() sources = set() for match in re.finditer(r'(?i)art\.?\s*\d+', answer): article_ref = match.group(0).strip() for idx, source in self.model.source_mapper.idx_to_source.items(): if article_ref.lower() in source.lower(): sources.add(source) return {"answer": answer, "sources": list(sources)} def main(): source_mapper = SourceMapper() model_name = "crumb/nano-mistral" tokenizer = AutoTokenizer.from_pretrained(model_name) tokenizer.pad_token = tokenizer.eos_token # Przygotowanie danych catalog_path = "file_catalog.json" data = prepare_dataset("files", catalog_path, source_mapper) if not data: print("\nBrak danych do treningu!") return dataset = Dataset.from_list(data) def tokenize(examples): return tokenizer( examples["text"], truncation=True, padding="max_length", max_length=512, return_tensors="pt" ) tokenized_dataset = dataset.map(tokenize, batched=True, batch_size=16) training_args = TrainingArguments( output_dir="./results", num_train_epochs=8, per_device_train_batch_size=4, gradient_accumulation_steps=8, learning_rate=5e-6, weight_decay=0.01, warmup_ratio=0.1, fp16=torch.cuda.is_available(), logging_steps=50, save_strategy="epoch", eval_strategy="no", report_to="none", remove_unused_columns=False ) model = CustomModel(model_name, AutoModelForCausalLM.from_pretrained(model_name).config) model.source_mapper = source_mapper model.to("cuda" if torch.cuda.is_available() else "cpu") trainer = CustomTrainer( model=model, args=training_args, train_dataset=tokenized_dataset, tokenizer=tokenizer ) print("\nRozpoczęcie treningu...") trainer.train() print("\nKońcowa ewaluacja...") trainer.evaluate() if __name__ == "__main__": main()