diff --git a/gemma.py b/gemma.py new file mode 100644 index 0000000..5742741 --- /dev/null +++ b/gemma.py @@ -0,0 +1,75 @@ +import torch +import chromadb +from sentence_transformers import SentenceTransformer +from datasets import Dataset +from peft import LoraConfig, get_peft_model +from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer + +# 1️⃣ Inicjalizacja ChromaDB i modelu do embeddingów +chroma_client = chromadb.PersistentClient(path="./chroma_db") +collection = chroma_client.get_or_create_collection("my_embeddings") +embed_model = SentenceTransformer("all-MiniLM-L6-v2") + +# 2️⃣ Dodanie dokumentów i embeddingów +documents = [ + "Jak założyć firmę w Polsce?", + "Jak rozliczyć podatek VAT?", + "Procedura składania reklamacji w e-sklepie.", + "Jakie dokumenty są potrzebne do rejestracji działalności?" +] +embeddings = embed_model.encode(documents).tolist() + +for i, (doc, emb) in enumerate(zip(documents, embeddings)): + collection.add(ids=[str(i)], documents=[doc], embeddings=[emb]) + +# 3️⃣ Przygotowanie danych treningowych +def create_training_data(): + data = collection.get(include=["documents", "embeddings"]) + return Dataset.from_dict({ + "text": data["documents"], + "embedding": data["embeddings"] + }) + +dataset = create_training_data() + +# 4️⃣ Ładowanie modelu Gemma 2 7B +device = "cuda" if torch.cuda.is_available() else "cpu" +model_name = "google/gemma-7b" +model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16).to(device) +tokenizer = AutoTokenizer.from_pretrained(model_name) + +# 5️⃣ Konfiguracja LoRA dla efektywnego treningu +lora_config = LoraConfig( + r=8, lora_alpha=32, lora_dropout=0.1, bias="none", task_type="CAUSAL_LM" +) +model = get_peft_model(model, lora_config) + +# 6️⃣ Tokenizacja danych +def tokenize_function(examples): + return tokenizer(examples["text"], padding="max_length", truncation=True) + +tokenized_dataset = dataset.map(tokenize_function, batched=True) + +# 7️⃣ Parametry treningu +training_args = TrainingArguments( + output_dir="./results", + per_device_train_batch_size=2, + num_train_epochs=3, + logging_dir="./logs", + save_strategy="epoch" +) + +# 8️⃣ Trening modelu +trainer = Trainer( + model=model, + args=training_args, + train_dataset=tokenized_dataset, +) + +trainer.train() + +# 9️⃣ Zapisanie dostrojonego modelu +model.save_pretrained("./trained_model/gemma") +tokenizer.save_pretrained("./trained_model/gemma") + +print("✅ Model został wytrenowany i zapisany!")