Update app.py

app.py

@@ -8,6 +8,9 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments,
 from peft import LoraConfig, get_peft_model
 import gradio as gr
 from transformers import pipeline
+from unsloth import FastLanguageModel, is_bfloat16_supported
+import torch
+from trl import SFTTrainer
 
 
 def extract_text_from_pdf(pdf_path):
@@ -36,7 +39,13 @@ login(token=hf_token)
 #model_name = "meta-llama/Llama-2-7b-hf"  # You can use a smaller one like "meta-llama/Llama-2-7b-chat-hf"
 model_name = "unsloth/llama-2-7b-chat"
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+#tokenizer = AutoTokenizer.from_pretrained(model_name)
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name=model_name,
+    max_seq_length=2048
+)
+
+model = FastLanguageModel.get_peft_model(model)
 
 # Create dataset
 data = {"text": [clean_text]}
@@ -55,43 +64,31 @@ def tokenize_function(examples):
 
 tokenized_datasets = dataset.map(tokenize_function, batched=True)
 
-# Load LLaMA 2 model in 4-bit mode to save memory
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    load_in_4bit=True,  # Use 4-bit quantization for efficiency
-    device_map="auto"
-    #device_map="cpu",
-    #quantization_config=None
-)
-
-# Apply LoRA (efficient fine-tuning)
-lora_config = LoraConfig(
-    r=8,  # Low-rank parameter
-    lora_alpha=32,
-    target_modules=["q_proj", "v_proj"],  # Applies only to attention layers
-    lora_dropout=0.05
-)
-
-model = get_peft_model(model, lora_config)
-
-training_args = TrainingArguments(
-    output_dir="./results",
-    evaluation_strategy="no",   # Disable evaluation (to enable, change value to 'epoch')
-    learning_rate=2e-4,
-    per_device_train_batch_size=1,  # Reduce batch size for memory efficiency
-    per_device_eval_batch_size=1,
-    num_train_epochs=3,
-    weight_decay=0.01,
-    save_strategy="epoch",
-    logging_dir="./logs",
-    logging_steps=10,
-)
-
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized_datasets,
-    tokenizer=tokenizer,
+trainer = SFTTrainer(
+    model = model,
+    tokenizer = tokenizer,
+    train_dataset = dataset,
+    dataset_text_field = "text",
+    max_seq_length = max_seq_length,
+    dataset_num_proc = 2,
+    packing = False, # Can make training 5x faster for short sequences.
+    args = TrainingArguments(
+        per_device_train_batch_size = 2,
+        gradient_accumulation_steps = 4,
+        warmup_steps = 5,
+        # num_train_epochs = 1, # Set this for 1 full training run.
+        max_steps = 60,
+        learning_rate = 2e-4,
+        fp16 = not is_bfloat16_supported(),
+        bf16 = is_bfloat16_supported(),
+        logging_steps = 1,
+        optim = "adamw_8bit",
+        weight_decay = 0.01,
+        lr_scheduler_type = "linear",
+        seed = 3407,
+        output_dir = "outputs",
+        report_to = "none", # Use this for WandB etc
+    ),
 )
 
 def perform_training():