File size: 3,346 Bytes
b05b1d8 |
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 |
# Importing necessary libraries
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments, TextClassificationPipeline
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
import gradio as gr
# Load the dataset
ds = load_dataset("GonzaloA/fake_news")
# Load pre-trained tokenizer
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
# Define tokenization function
def tokenize_function(examples):
return tokenizer(examples['text'], padding='max_length', truncation=True, max_length=128)
# Apply tokenization
tokenized_datasets = ds.map(tokenize_function, batched=True)
# Load pre-trained BERT model for sequence classification
model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
# Define training arguments
training_args = TrainingArguments(
output_dir='./results',
num_train_epochs=3,
per_device_train_batch_size=8,
per_device_eval_batch_size=8,
evaluation_strategy='epoch',
logging_dir='./logs',
)
# Create trainer instance
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_datasets['train'].shuffle().select(range(1000)),
eval_dataset=tokenized_datasets['test'].shuffle().select(range(1000)),
)
# Start training
trainer.train()
# Define function to compute metrics
def compute_metrics(pred):
labels = pred.label_ids
preds = pred.predictions.argmax(-1)
precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='binary')
acc = accuracy_score(labels, preds)
return {'accuracy': acc, 'f1': f1, 'precision': precision, 'recall': recall}
# Update trainer to include custom metrics
trainer.compute_metrics = compute_metrics
# Evaluate the model
eval_result = trainer.evaluate()
print(eval_result)
# Save the fine-tuned model and tokenizer
trainer.save_model('TeamQuad-fine-tuned-bert')
tokenizer.save_pretrained('TeamQuad-fine-tuned-bert')
# Load the fine-tuned model and tokenizer
new_model = AutoModelForSequenceClassification.from_pretrained('TeamQuad-fine-tuned-bert')
new_tokenizer = AutoTokenizer.from_pretrained('TeamQuad-fine-tuned-bert')
# Create a classification pipeline
classifier = TextClassificationPipeline(model=new_model, tokenizer=new_tokenizer)
# Add label mapping for fake news detection (assuming LABEL_0 = 'fake' and LABEL_1 = 'true')
label_mapping = {0: 'fake', 1: 'true'}
# Function to classify input text
def classify_news(text):
result = classifier(text)
# Extract the label and score
label = result[0]['label'] # 'LABEL_0' or 'LABEL_1'
score = result[0]['score'] # Confidence score
mapped_result = {'label': label_mapping[int(label.split('_')[1])], 'score': score}
return f"Label: {mapped_result['label']}, Score: {mapped_result['score']:.4f}"
# Create a Gradio interface
iface = gr.Interface(
fn=classify_news, # The function to process the input
inputs=gr.Textbox(lines=10, placeholder="Enter a news headline or article to classify..."),
outputs="text", # Output will be displayed as text
title="Fake News Detection",
description="Enter a news headline or article and see whether the model classifies it as 'Fake News' or 'True News'.",
)
# Launch the interface
iface.launch(share=True)
|