jonACE/model-finetuning-with-own-data

1import fitz  # PyMuPDF for PDF extraction2import re3import unsloth4import os5from huggingface_hub import login6from datasets import Dataset7from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer8from peft import LoraConfig, get_peft_model9import gradio as gr10from transformers import pipeline11 12 13def extract_text_from_pdf(pdf_path):14    """Extract text from a PDF file"""15    doc = fitz.open(pdf_path)16    text = "\n".join([page.get_text("text") for page in doc])17    return text.strip()18 19def preprocess_text(text):20    """Basic text preprocessing"""21    return re.sub(r"\s+", " ", text).strip()22 23pdf_text = extract_text_from_pdf("new-american-standard-bible.pdf")24clean_text = preprocess_text(pdf_text)25 26 27# Read the Hugging Face token from environment variables28hf_token = os.getenv("access_token")29 30if hf_token is None:31    raise ValueError("'access_token' is not set. Add it as a secret variable in Hugging Face Spaces.")32 33# Log in to Hugging Face34login(token=hf_token)35 36#model_name = "meta-llama/Llama-2-7b-hf"  # You can use a smaller one like "meta-llama/Llama-2-7b-chat-hf"37model_name = "unsloth/llama-2-7b-chat"38 39tokenizer = AutoTokenizer.from_pretrained(model_name)40 41# Create dataset42data = {"text": [clean_text]}43dataset = Dataset.from_dict(data)44 45# Set a padding token manually46tokenizer.pad_token = tokenizer.eos_token  # Use EOS as PAD token47# Alternatively, add a new custom pad token48# tokenizer.add_special_tokens({'pad_token': '[PAD]'})49 50# Tokenization function51def tokenize_function(examples):52    tokens = tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)53    tokens["labels"] = tokens["input_ids"].copy()  # Use input as labels for text generation54    return tokens55 56tokenized_datasets = dataset.map(tokenize_function, batched=True)57 58# Load LLaMA 2 model in 4-bit mode to save memory59model = AutoModelForCausalLM.from_pretrained(60    model_name,61    load_in_4bit=True,  # Use 4-bit quantization for efficiency62    device_map="auto"63    #device_map="cpu",64    #quantization_config=None65)66 67# Apply LoRA (efficient fine-tuning)68lora_config = LoraConfig(69    r=8,  # Low-rank parameter70    lora_alpha=32,71    target_modules=["q_proj", "v_proj"],  # Applies only to attention layers72    lora_dropout=0.0573)74 75model = get_peft_model(model, lora_config)76 77training_args = TrainingArguments(78    output_dir="./results",79    evaluation_strategy="no",   # Disable evaluation (to enable, change value to 'epoch')80    learning_rate=2e-4,81    per_device_train_batch_size=1,  # Reduce batch size for memory efficiency82    per_device_eval_batch_size=1,83    num_train_epochs=3,84    weight_decay=0.01,85    save_strategy="epoch",86    logging_dir="./logs",87    logging_steps=10,88)89 90trainer = Trainer(91    model=model,92    args=training_args,93    train_dataset=tokenized_datasets,94    tokenizer=tokenizer,95)96 97def perform_training():98    trainer.train()99 100perform_training()101 102model.save_pretrained("./fine_tuned_llama2")103tokenizer.save_pretrained("./fine_tuned_llama2")104 105 106# CHATBOT START107chatbot = pipeline("text-generation", model="./fine_tuned_llama2")108 109def chatbot_response(prompt):110    result = chatbot(prompt, max_length=100, do_sample=True, temperature=0.7)111    return result[0]["generated_text"]112 113iface = gr.Interface(fn=chatbot_response, inputs="text", outputs="text")114iface.launch()