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


d
examples = [
    "Machine Learning is at the forefront of education, replacing human jobs",
    "AI model leaves scientists confused",
    "This model is not really cool",
]

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