import json
from pathlib import Path

import gradio as gr
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TextClassificationPipeline,
)

labels = [
    # 'agency',
    # 'humanComparison',
    # 'hyperbole',
    # 'historyComparison',
    # 'unjustClaims',
    # 'deepSounding',
    # 'skeptics',
    # 'deEmphasize',
    # 'performanceNumber',
    # 'inscrutable',
    # 'objective'
    "agency",
    # "suggestiveImagery",
    "comparisonWithHumanIntelligence",
    "comparisonWithHumanSkills",
    "hyperbole",
    "uncriticalHistoryComparison",
    "unjustifiedClaimsAboutFuture",
    "falseClaimsAboutProgress",
    "incorrectClaimsAboutStudyReport",
    "deepSoundingTermsForBanalities",
    "treatingSpokespeopleAsNeutral",
    "repeatingPRTerms",
    "noDiscussionOfLimitations",
    "deEmphasizingLimitations",
    "limitationsAddressedBySkeptics",
    "downplayingHumanLabour",
    "performanceNumbersWithoutCaveats",
    # "inscrutability",
]

models = {}

pipes = {}

tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

for label in labels:
    models[label] = AutoModelForSequenceClassification.from_pretrained(
        f'xt0r3/aihype_{label}-vs-rest',
    )
    pipes[label] = TextClassificationPipeline(
        model=models[label], tokenizer=tokenizer, top_k=None,
    )


def predict(text):
    preds = {}

    for label in labels:
        pred_array = pipes[label](text)[0]
        for pred in pred_array:
            if pred['label'] == 'LABEL_1':
                preds[label] = pred['score']

    return preds


examples = [
    "Machine Learning is at the forefront of education, replacing human jobs",
    "Machine Learning is at the forefront of education, but it isn't replacing human jobs just yet",
    "AI model leaves scientists confused",
    "Tech company unveils radical new assistant that will 'do everything for you'",
]

intf = gr.Interface(fn=predict, inputs="textbox",
                    outputs="label", examples=examples)
intf.launch(inline=False)