import json
import pandas as pd

def preprocess_json(data, possible_relationships):
    # Extract the relevant information
    document = data['document']
    tokens = data['tokens']
    relations = data['relations']

    # Create a mapping of token index to its text and entity label
    token_info = {i: {'text': t['text'], 'label':     t['entityLabel']} for i, t in enumerate(tokens)}
    
# Format the entities and relationships
    entities = [(t['text'], t['entityLabel']) for t in tokens]
    formatted_entities = ", ".join([f"{text} ({label})" for text, label in entities])

    formatted_relations = []
    for r in relations:
        child_index = r['child']
        head_index = r['head']

        if child_index < len(tokens) and head_index < len(tokens):
            child = token_info[child_index]['text']
            head = token_info[head_index]['text']
            
 relation_label = r['relationLabel']
            formatted_relations.append(f"{child} -> {head} ({relation_label})")

    formatted_relations = "; ".join(formatted_relations)

    # Create the formatted prompt and response
    prompt = f"systemExtract relationships between entities from the following text.user Text: \"{document}\" Entities: {formatted_entities}. Possible relationships: {', '.join(possible_relationships)}."
    response = f"assistantThe relations between the entities: {formatted_relations}"
    full_prompt = prompt + response
    return full_prompt

# List of possible relationships (customize as needed)
possible_relationships = ["MUST_HAVE", "REQUIRES", "NICE_TO_HAVE"]

input_path = "/content/UBIAI_REL_data.json"
# Read the input JSON file
with open(input_path, 'r') as file:
     data = json.load(file)

# Preprocess all JSON strings
data = [preprocess_json(j, possible_relationships) for j in data]

# Convert to a DataFrame
df = pd.DataFrame(data, columns=["text"])

# Save to CSV
df.to_csv('fine_tuning_data.csv', index=False)


				
			

				
					import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import prepare_model_for_kbit_training, LoraConfig
from trl import setup_chat_format

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

# Tokenizer setup
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True, trust_remote_code=True)

tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.eos_token_id
tokenizer.padding_side = 'left'
tokenizer.model_max_length = 2048

# Quantization configuration
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16
)

# Model setup
device_map = {"": torch.cuda.current_device()} if torch.cuda.is_available() else None
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map=device_map,
    quantization_config=bnb_config
)

model, tokenizer = setup_chat_format(model, tokenizer)
model = prepare_model_for_kbit_training(model)

# LoRA configuration

peft_config = LoraConfig(
    lora_alpha=128,
    lora_dropout=0.05,
    r=256,
    bias="none",
    target_modules=["q_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "k_proj", "v_proj"],
    task_type="CAUSAL_LM",)