| # 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() | |