File size: 2,588 Bytes
39a8b03 54d96e0 2a393b0 39a8b03 40f0be0 39a8b03 40f0be0 39a8b03 40f0be0 39a8b03 2a393b0 39a8b03 54d96e0 39a8b03 54d96e0 39a8b03 40f0be0 39a8b03 |
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 |
import gradio as gr
import json
from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments, AutoTokenizer
from datasets import Dataset
import shutil
import torch
# Load tokenized data
def load_data(tokenized_file):
with open(tokenized_file.name, 'r') as f:
tokenized_data = json.load(f)
return tokenized_data
# Fine-tune the model
def fine_tune_model(tokenized_file, progress=gr.Progress()):
tokenized_data = load_data(tokenized_file)
# Convert tokenized data to Dataset
dataset = Dataset.from_dict(tokenized_data)
# Split the dataset into train and validation sets
tokenized_datasets = dataset.train_test_split(test_size=0.2)
model = AutoModelForSequenceClassification.from_pretrained('anferico/bert-for-patents', num_labels=2)
tokenizer = AutoTokenizer.from_pretrained('anferico/bert-for-patents')
training_args = TrainingArguments(
output_dir='./results',
num_train_epochs=3,
per_device_train_batch_size=8, # Reduce batch size
per_device_eval_batch_size=16,
gradient_accumulation_steps=2, # Use gradient accumulation
fp16=True, # Enable mixed precision training
warmup_steps=500,
weight_decay=0.01,
logging_dir='./logs',
logging_steps=10,
evaluation_strategy="epoch",
save_strategy="epoch",
load_best_model_at_end=True,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_datasets['train'],
eval_dataset=tokenized_datasets['test'],
)
progress(0.5, "Fine-tuning the model...")
trainer.train()
progress(1.0, "Fine-tuning complete.")
model.save_pretrained('./fine_tuned_patentbert')
tokenizer.save_pretrained('./fine_tuned_patentbert')
# Create a zip file of the fine-tuned model
shutil.make_archive('fine_tuned_patentbert', 'zip', './fine_tuned_patentbert')
return "Model fine-tuned and saved successfully. Download the model using the link below.", "fine_tuned_patentbert.zip"
# Create Gradio interface
iface = gr.Interface(
fn=fine_tune_model,
inputs=[
gr.File(label="Upload Tokenized Data JSON")
],
outputs=[
gr.Textbox(label="Processing Information"),
gr.File(label="Download Fine-Tuned Model")
],
title="Fine-Tune Patent BERT Model",
description="Upload tokenized JSON file to fine-tune the BERT model.",
live=True # Enable live updates for progress
)
# Launch the interface
iface.launch()
|