import os import numpy as np import pandas as pd import os from tqdm import tqdm from transformers import pipeline from transformers import AutoTokenizer, FalconForCausalLM import torch from datasets import Dataset from peft import LoraConfig from trl import SFTTrainer from transformers import ( AutoTokenizer, BitsAndBytesConfig, TrainingArguments, pipeline, ) from sklearn.metrics import (accuracy_score, classification_report, confusion_matrix) from sklearn.model_selection import train_test_split os.environ["CUDA_VISIBLE_DEVICES"] = "0" os.environ["TOKENIZERS_PARALLELISM"] = "false" import warnings warnings.filterwarnings("ignore") def generate_prompt(data_point): return f"""### Instruction: Classify whether the given chunk involves a decision that will effect the story or not. A decision is defined as when the character goes about making a choice between two or more options. The decision should be significant enough to affect the story in a major way. It doesn't really involve emotions, feelings or thoughts, but what the character does, or what happens to them. This involes interactions between characters, or the character and the environment. What isn't a decision is chunks describing the setting, or the character's thoughts or feelings. Return the answer as the corresponding decision label "yes" or "no" ### Text: {data_point["text"]} ### Decision: {data_point["decision"]} """ def generate_test_prompt(data_point): return f"""### Instruction: Classify whether the given chunk involves a decision that will effect the story or not. A decision is defined as when the character goes about making a choice between two or more options. The decision should be significant enough to affect the story in a major way. It doesn't really involve emotions, feelings or thoughts, but what the character does, or what happens to them. This involes interactions between characters, or the character and the environment. What isn't a decision is chunks describing the setting, or the character's thoughts or feelings. Return the answer as the corresponding decision label "yes" or "no" ### Text: {data_point["text"]} ### Decision: """ def predict(X_test, model, tokenizer): y_pred = [] for i in tqdm(range(len(X_test))): prompt = X_test.iloc[i]["text"] pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_new_tokens = 1, temperature = 0.0, ) result = pipe(prompt, pad_token_id=pipe.tokenizer.eos_token_id) answer = result[0]['generated_text'].split("=")[-1].lower() if "yes" in answer: y_pred.append("yes") elif "no" in answer: y_pred.append("no") else: y_pred.append("none") return y_pred def evaluate(y_true, y_pred): labels = ['yes', 'no', 'none'] mapping = {"yes": 1, "no": 0, 'none':2} def map_func(x): return mapping.get(x, 1) y_true = np.vectorize(map_func)(y_true) y_pred = np.vectorize(map_func)(y_pred) # Calculate accuracy accuracy = accuracy_score(y_true=y_true, y_pred=y_pred) print(f'Accuracy: {accuracy:.3f}') # Generate accuracy report unique_labels = set(y_true) # Get unique labels for label in unique_labels: label_indices = [i for i in range(len(y_true)) if y_true[i] == label] label_y_true = [y_true[i] for i in label_indices] label_y_pred = [y_pred[i] for i in label_indices] accuracy = accuracy_score(label_y_true, label_y_pred) print(f'Accuracy for label {label}: {accuracy:.3f}') # Generate classification report class_report = classification_report(y_true=y_true, y_pred=y_pred) print('\nClassification Report:') print(class_report) # Generate confusion matrix conf_matrix = confusion_matrix(y_true=y_true, y_pred=y_pred, labels=[0, 1, 2]) print('\nConfusion Matrix:') print(conf_matrix) def prep_data(): filename = '../../data/output/decisions.csv' df = pd.read_csv(filename, encoding="utf-8", encoding_errors="replace") df = df[['text', 'decision']] X_train = list() X_test = list() for decision in ["yes", "no"]: train, test = train_test_split(df[df.decision==decision], train_size=.8, test_size=.2, random_state=42) X_train.append(train) X_test.append(test) X_train = pd.concat(X_train).sample(frac=1, random_state=10) X_test = pd.concat(X_test) eval_idx = [idx for idx in df.index if idx not in list(train.index) + list(test.index)] X_eval = df[df.index.isin(eval_idx)] X_eval = (X_eval .groupby('decision', group_keys=False) .apply(lambda x: x.sample(n=50, random_state=10, replace=True))) X_train = X_train.reset_index(drop=True) X_train = pd.DataFrame(X_train.apply(generate_prompt, axis=1), columns=["text"]) X_eval = pd.DataFrame(X_eval.apply(generate_prompt, axis=1), columns=["text"]) y_true = X_test.decision X_test = pd.DataFrame(X_test.apply(generate_test_prompt, axis=1), columns=["text"]) train_data = Dataset.from_pandas(X_train) eval_data = Dataset.from_pandas(X_eval) return train_data, eval_data def prep_model(): model_name = "Rocketknight1/falcon-rw-1b" compute_dtype = getattr(torch, "float16") bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_use_double_quant=False, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=compute_dtype, ) model = FalconForCausalLM.from_pretrained( model_name, device_map="auto", quantization_config=bnb_config, ) model.config.use_cache = False model.config.pretraining_tp = 1 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left", add_bos_token=True, add_eos_token=True, ) tokenizer.pad_token = tokenizer.eos_token return model, tokenizer def prep_trainer(): OUTPUT_DIR = "falcon-clf" train_data, eval_data = prep_data() model, tokenizer = prep_model() peft_config = LoraConfig( lora_alpha=16, lora_dropout=0.1, r=64, bias="none", task_type="CAUSAL_LM", ) training_arguments = TrainingArguments( output_dir=OUTPUT_DIR, num_train_epochs=20, per_device_train_batch_size=1, gradient_accumulation_steps=8, # 4 optim="paged_adamw_32bit", save_steps=0, logging_steps=10, learning_rate=2e-4, weight_decay=0.001, fp16=True, bf16=False, max_grad_norm=0.3, max_steps=-1, warmup_ratio=0.03, group_by_length=True, lr_scheduler_type="cosine", report_to="tensorboard", evaluation_strategy="epoch" ) trainer = SFTTrainer( model=model, train_dataset=train_data, eval_dataset=eval_data, peft_config=peft_config, dataset_text_field="text", tokenizer=tokenizer, args=training_arguments, packing=False, max_seq_length=1024, ) return trainer def train_model(): trainer = prep_trainer() trainer.train() trainer.model.save_pretrained("falcon-clf") trainer.push_to_hub() def get_classifier(): classifier = pipeline(model=f"suneeln-duke/falcon-clf", device_map="auto") return classifier def classify_dec(text, classifier): text = generate_test_prompt({ 'text': text }) result = classifier(text, pad_token_id=classifier.tokenizer.eos_token_id) answer = result[0]['generated_text'].split("=")[-1].lower() if "yes" in answer: return "yes" elif "no" in answer: return "no"