File size: 2,154 Bytes
a30b42a |
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 |
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
from datasets import load_dataset
# Define datasets and their IDs
datasets_info = {
"SQuAD": "squad",
"SQuAD 2.0": "squad_v2",
"Natural Questions": "nq",
"TriviaQA": "triviaqa",
"QuAC": "quac",
"FAQ Dataset": "faq",
"BoolQ": "boolq",
"Open Book QA": "obqa"
}
# Load model and tokenizer directly
tokenizer = AutoTokenizer.from_pretrained("nvidia/Llama-3.1-Nemotron-70B-Instruct-HF")
model = AutoModelForCausalLM.from_pretrained("nvidia/Llama-3.1-Nemotron-70B-Instruct-HF")
def train_model(dataset_name):
# Load the dataset
dataset = load_dataset(datasets_info[dataset_name])
# Tokenization
def preprocess_function(examples):
return tokenizer(examples['question'], examples['context'], truncation=True)
tokenized_dataset = dataset.map(preprocess_function, batched=True)
# Fine-tune the model
training_args = TrainingArguments(
output_dir=f"./{dataset_name}_model",
evaluation_strategy="epoch",
learning_rate=2e-5,
per_device_train_batch_size=8,
per_device_eval_batch_size=8,
num_train_epochs=3,
weight_decay=0.01,
logging_dir='./logs',
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_dataset['train'],
eval_dataset=tokenized_dataset['validation']
)
trainer.train()
# Save the model weights
model.save_pretrained(f"./{dataset_name}_model")
tokenizer.save_pretrained(f"./{dataset_name}_model")
return f"Model trained and saved for {dataset_name}!"
# Gradio Interface
with gr.Blocks() as demo:
gr.Markdown("## Train QA Model on Multiple Datasets")
dataset_name = gr.Dropdown(choices=list(datasets_info.keys()), label="Select Dataset")
train_button = gr.Button("Train Model")
output = gr.Textbox(label="Output")
def train_and_display(dataset_name):
return train_model(dataset_name)
train_button.click(train_and_display, inputs=dataset_name, outputs=output)
demo.launch()
|