File size: 8,855 Bytes
6c57304
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
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"