!pip install -q groq!pip install -U accelerate bitsandbytes datasets evaluate!pip install -U peft transformers trl

# For Google Colab settingsfrom google.colab import userdata, drive
# This will prompt for authorizationdrive.mount('/content/drive')
# clone project! git clone https://github.com/mcks2000/llm_notebooks.git

# Set the working directory%cd '/content/llm_notebooks/knowlage_graph/llama3_re'

# For Hugging Face Hub settingfrom huggingface_hub import login
# Upload the HuggingFace token (should have WRITE access) from Colab secretsHF = userdata.get('HF')
# This is needed to upload the model to HuggingFacelogin(token=HF,add_to_git_credential=True)

# Create a path variable for the data folderdata_path = '/content/llm_notebooks/knowlage_graph/llama3_re/datas/'
# Full fine-tuning datasetsft_dataset_file = f'{data_path}sft_train_data.json'
# Data collected from the the mini-testmini_data_path = f'{data_path}mini_data.json'
# Test data containing all three outputsall_tests_data = f'{data_path}all_tests.json'
# The adjusted training datasettrain_data_path = f'{data_path}sft_train_data.json'
# Create a path variable for the SFT model to be saved locallysft_model_path = '/content/llm_notebooks/knowlage_graph/llama3_re/Llama3_RE/'

from datasets import load_dataset
# Load the datasetdataset = load_dataset("databricks/databricks-dolly-15k")
# Choose the desired category from the datasetie_category = [e for e in dataset["train"] if e["category"]=="information_extraction"]
# Retain only the context from each instanceie_context = [e["context"] for e in ie_category]
# Split the text into sentences (at the period) and keep the first sentencereduced_context = [text.split('.')[0] + '.' for text in ie_context]
# Retain sequences of specified lengths only (use character length)sampler = [e for e in reduced_context if 30 < len(e) < 170]

system_message = """You are an experienced annontator.Extract all entities and the relations between them from the following text.Write the answer as a triple entity1|relationship|entitity2.Do not add anything else.Example Text: Alice is from France.Answer: Alice|is from|France."""

messages = [[{"role": "system","content": f"{system_message}"},{"role": "user", "content": e}] for e in sampler]

import osfrom groq import Groq
gclient = Groq(api_key=userdata.get("GROQ"),)

import timefrom tqdm import tqdm
def process_data(prompt):
"""Send one request and retrieve model's generation."""
chat_completion = gclient.chat.completions.create(messages=prompt, # input prompt to send to the modelmodel="llama3-70b-8192", # according to GroqCloud labelingtemperature=0.5, # controls diversitymax_tokens=128, # max number tokens to generatetop_p=1, # proportion of likelihood weighted options to considerstop=None, # string that signals to stop generatingstream=False, # if set partial messages are sent)return chat_completion.choices[0].message.content
def send_messages(messages):
"""Process messages in batches with a pause between batches."""
 batch_size = 10answers = []
for i in tqdm(range(0, len(messages), batch_size)): # batches of size 10
batch = messages[i:i+10]# get the next batch of messages
for message in batch:output = process_data(message)answers.append(output)
if i + 10 < len(messages):# check if there are batches lefttime.sleep(10)# wait for 10 seconds
return answers

# Data generation with Llama3-70Banswers = send_messages(messages)
# Combine input data with the generated datasetcombined_dataset = [{'text': user, 'gold_re': output} for user, output in zip(sampler, answers)]

import randomrandom.seed(17)
# Select 20 random entriesmini_data = random.sample(combined_dataset, 20)
# Build conversational formatparsed_mini_data = [[{'role': 'system', 'content': system_message}, {'role': 'user', 'content': e['text']}] for e in mini_data]
# Create the training settrain_data = [item for item in combined_dataset if item not in mini_data]

{'text': 'Long before any knowledge of electricity existed, people were aware of shocks from electric fish.', 'gold_re': 'people|were aware of|shocks\\nshocks|from|electric fish\\nelectric fish|had|electricity', 'test_re': 'electric fish|were aware of|shocks'}

def create_conversation(sample):return {"messages": [{"role": "system","content": system_message},{"role": "user", "content": sample["text"]},{"role": "assistant", "content": sample["gold_re"]}]}
from datasets import load_dataset, Dataset
train_dataset = Dataset.from_list(train_data)
# Transform to conversational formattrain_dataset = train_dataset.map(create_conversation,remove_columns=train_dataset.features,batched=False)

model_id="meta-llama/Meta-Llama-3-8B"

from transformers import AutoTokenizer
# Tokenizertokenizer = AutoTokenizer.from_pretrained(model_id,use_fast=True,trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_tokentokenizer.pad_token_id =tokenizer.eos_token_idtokenizer.padding_side = 'left'
# Set a maximum lengthtokenizer.model_max_length = 512

from transformers import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_use_double_quant=True,bnb_4bit_quant_type="nf4",bnb_4bit_compute_dtype=torch.bfloat16)

from transformers import AutoModelForCausalLMfrom peft import prepare_model_for_kbit_trainingfrom trl import setup_chat_format
device_map = {"": torch.cuda.current_device()} if torch.cuda.is_available() else None
model = AutoModelForCausalLM.from_pretrained(model_id,device_map=device_map,attn_implementation="flash_attention_2",quantization_config=bnb_config)
model, tokenizer = setup_chat_format(model, tokenizer)model = prepare_model_for_kbit_training(model)

from peft import LoraConfig
# According to Sebastian Raschka findingspeft_config = LoraConfig(lora_alpha=128, #32lora_dropout=0.05,r=256,#16bias="none",target_modules=["q_proj", "o_proj", "gate_proj", "up_proj","down_proj", "k_proj", "v_proj"],task_type="CAUSAL_LM",)

from transformers import TrainingArguments
# Adapted fromPhil Schmid blogpostargs = TrainingArguments(output_dir=sft_model_path,# directory to save the model and repository idnum_train_epochs=2, # number of training epochsper_device_train_batch_size=4,# batch size per device during traininggradient_accumulation_steps=2,# number of steps before performing a backward/update passgradient_checkpointing=True,# use gradient checkpointing to save memory, use in distributed trainingoptim="adamw_8bit", # choose paged_adamw_8bit if not enough memorylogging_steps=10, # log every 10 stepssave_strategy="epoch",# save checkpoint every epochlearning_rate=2e-4, # learning rate, based on QLoRA paperbf16=True,# use bfloat16 precisiontf32=True,# use tf32 precisionmax_grad_norm=0.3,# max gradient norm based on QLoRA paperwarmup_ratio=0.03,# warmup ratio based on QLoRA paperlr_scheduler_type="constant", # use constant learning rate schedulerpush_to_hub=True, # push model to Hugging Face hubhub_model_id="llama3-8b-sft-qlora-re",report_to="tensorboard", # report metrics to tensorboard)

from trl import SFTTrainer
trainer = SFTTrainer(model=model,args=args,train_dataset=sft_dataset,peft_config=peft_config,max_seq_length=512,tokenizer=tokenizer,packing=False, # True if the dataset is largedataset_kwargs={"add_special_tokens": False,# the template adds the special tokens"append_concat_token": False, # no need to add additional separator token})
trainer.train()trainer.save_model()

import torchimport gcdel modeldel tokenizergc.collect()torch.cuda.empty_cache()

from peft import AutoPeftModelForCausalLMfrom transformers import AutoTokenizer, pipelineimport torch
# HF modelpeft_model_id = "solanaO/llama3-8b-sft-qlora-re"
# Load Model with PEFT adaptermodel = AutoPeftModelForCausalLM.from_pretrained(peft_model_id,device_map="auto",torch_dtype=torch.float16,offload_buffers=True)

okenizer = AutoTokenizer.from_pretrained(peft_model_id)
tokenizer.pad_token = tokenizer.eos_tokentokenizer.pad_token_id =tokenizer.eos_token_id

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

def create_input_prompt(sample):return {"messages": [{"role": "system","content": system_message},{"role": "user", "content": sample["text"]},]}
from datasets import Dataset
test_dataset = Dataset.from_list(mini_data)
# Transform to conversational formattest_dataset = test_dataset.map(create_input_prompt,remove_columns=test_dataset.features,batched=False)

# Generate the input promptprompt = pipe.tokenizer.apply_chat_template(test_dataset[2]["messages"][:2],tokenize=False,add_generation_prompt=True)# Generate the outputoutputs = pipe(prompt,max_new_tokens=128,do_sample=False,temperature=0.1,top_k=50,top_p=0.1,)# Display the resultsprint(f"Question: {test_dataset[2]['messages'][1]['content']}\\n")print(f"Gold-RE: {test_sampler[2]['gold_re']}\\n")print(f"LLama3-8B-RE: {test_sampler[2]['test_re']}\\n")print(f"SFT-Llama3-8B-RE: {outputs[0]['generated_text'][len(prompt):].strip()}")

Question: Long before any knowledge of electricity existed, people were aware of shocks from electric fish.
Gold-RE: people|were aware of|shocksshocks|from|electric fishelectric fish|had|electricity
LLama3-8B-RE: electric fish|were aware of|shocks
SFT-Llama3-8B-RE: people|were aware of|shocks shocks|from|electric fish