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)