test-123 / app.py
TeamQuad's picture
app.py
81c90f7 verified
from transformers import AutoModelForSequenceClassification, AutoTokenizer, TextClassificationPipeline
import gradio as gr
# 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 = {'LABEL_0': 'fake', 'LABEL_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[label], '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)