File size: 6,447 Bytes
7286eeb 4ec5643 03e8a0d 3f72ea7 4ec5643 7286eeb 4ec5643 03e8a0d 4ec5643 ecd989d 4ec5643 3f72ea7 4ec5643 3dca0ff 0e915a7 fb5819f ecd989d 0e915a7 ecd989d 0e915a7 ecd989d fb5819f ecd989d 0e915a7 4ec5643 ecd989d 03e8a0d 4ec5643 03e8a0d 4ec5643 3f72ea7 4ec5643 ecd989d d33fdc8 ecd989d d33fdc8 ecd989d 4ec5643 d33fdc8 bd92a3c fbec35f ecd989d bd92a3c ecd989d bd92a3c 4ec5643 03e8a0d ecd989d 4ec5643 ecd989d 4ec5643 ecd989d 03e8a0d 4ec5643 ecd989d 4ec5643 ecd989d 7286eeb | 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 | # import gradio as gr
# from datasets import load_dataset
# from transformers import LongformerTokenizer, LongformerForSequenceClassification, Trainer, TrainingArguments
# import torch
# # Load and preprocess dataset
# def load_and_preprocess_data():
# dataset = load_dataset('json', data_files='data.json')['train']
# tokenizer = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096')
# def preprocess_function(examples):
# inputs = {
# 'input_ids': [],
# 'attention_mask': [],
# 'labels': []
# }
# for example in examples['prompt']:
# # Tokenize input
# tokenized = tokenizer(
# example,
# truncation=True,
# padding='max_length',
# max_length=128 # Increase max_length to suit model needs
# )
# inputs['input_ids'].append(tokenized['input_ids'])
# inputs['attention_mask'].append(tokenized['attention_mask'])
# inputs['labels'].append(1 if "chosen" in example else 0) # Assign binary labels
# inputs['input_ids'] = torch.tensor(inputs['input_ids'], dtype=torch.long)
# inputs['attention_mask'] = torch.tensor(inputs['attention_mask'], dtype=torch.long)
# inputs['labels'] = torch.tensor(inputs['labels'], dtype=torch.long)
# return inputs
# dataset = dataset.map(preprocess_function, batched=True)
# dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
# return dataset, tokenizer
# # Define and train the model
# def train_model(dataset):
# model = LongformerForSequenceClassification.from_pretrained('allenai/longformer-base-4096', num_labels=2)
# training_args = TrainingArguments(
# output_dir='./results',
# evaluation_strategy="epoch",
# learning_rate=2e-5,
# per_device_train_batch_size=4,
# per_device_eval_batch_size=4,
# num_train_epochs=3,
# weight_decay=0.01,
# )
# class CustomTrainer(Trainer):
# def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
# labels = inputs.pop("labels")
# outputs = model(**inputs)
# logits = outputs.logits
# loss = torch.nn.CrossEntropyLoss()(logits.view(-1, model.config.num_labels), labels.view(-1))
# return (loss, outputs) if return_outputs else loss
# trainer = CustomTrainer(
# model=model,
# args=training_args,
# train_dataset=dataset,
# eval_dataset=dataset,
# )
# trainer.train()
# return model
# # Gradio Interface function to predict
# def predict(text):
# inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
# outputs = model(**inputs)
# prediction = torch.argmax(outputs.logits, dim=-1).item()
# return f"Prediction: {'chosen' if prediction == 1 else 'rejected'}"
# # Main script
# if __name__ == "__main__":
# dataset, tokenizer = load_and_preprocess_data()
# model = train_model(dataset)
# gr.Interface(fn=predict, inputs=gr.Textbox(), outputs="text").launch()
import gradio as gr
from datasets import load_dataset
from transformers import LongformerTokenizer, LongformerForSequenceClassification, Trainer, TrainingArguments
import torch
from torch.utils.data import DataLoader
# Load and preprocess dataset
def load_and_preprocess_data():
dataset = load_dataset('json', data_files='data.json')['train']
tokenizer = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096')
def preprocess_function(examples):
inputs = {
'input_ids': [],
'attention_mask': [],
'labels': []
}
for example in examples['prompt']:
# Tokenize input
tokenized = tokenizer(
example,
truncation=True,
padding='max_length',
max_length=128 # Increase max_length to suit model needs
)
inputs['input_ids'].append(tokenized['input_ids'])
inputs['attention_mask'].append(tokenized['attention_mask'])
inputs['labels'].append(1 if "chosen" in example else 0) # Assign binary labels
inputs['input_ids'] = torch.tensor(inputs['input_ids'], dtype=torch.long)
inputs['attention_mask'] = torch.tensor(inputs['attention_mask'], dtype=torch.long)
inputs['labels'] = torch.tensor(inputs['labels'], dtype=torch.long)
return inputs
dataset = dataset.map(preprocess_function, batched=True)
dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
return dataset, tokenizer
# Define and train the model
def train_model(dataset):
model = LongformerForSequenceClassification.from_pretrained('allenai/longformer-base-4096', num_labels=2)
training_args = TrainingArguments(
output_dir='./results',
evaluation_strategy="epoch",
learning_rate=2e-5,
per_device_train_batch_size=2, # Try reducing the batch size
per_device_eval_batch_size=2,
gradient_accumulation_steps=2, # Accumulate gradients over 2 steps
num_train_epochs=3,
weight_decay=0.01,
)
class CustomTrainer(Trainer):
def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
labels = inputs.pop("labels")
outputs = model(**inputs)
logits = outputs.logits
loss = torch.nn.CrossEntropyLoss()(logits.view(-1, model.config.num_labels), labels.view(-1))
return (loss, outputs) if return_outputs else loss
trainer = CustomTrainer(
model=model,
args=training_args,
train_dataset=dataset,
eval_dataset=dataset,
)
trainer.train()
return model
# Gradio Interface function to predict
def predict(text):
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
outputs = model(**inputs)
prediction = torch.argmax(outputs.logits, dim=-1).item()
return f"Prediction: {'chosen' if prediction == 1 else 'rejected'}"
# Main script
if __name__ == "__main__":
dataset, tokenizer = load_and_preprocess_data()
model = train_model(dataset)
gr.Interface(fn=predict, inputs=gr.Textbox(), outputs="text").launch()
|