import os
PATH = '/data/' # at least 150GB storage needs to be attached
os.environ['TRANSFORMERS_CACHE'] = PATH
os.environ['HF_HOME'] = PATH
os.environ['HF_DATASETS_CACHE'] = PATH
os.environ['TORCH_HOME'] = PATH

import gradio as gr

from interfaces.cap import demo as cap_demo
from interfaces.manifesto import demo as manifesto_demo
from interfaces.sentiment import demo as sentiment_demo
from interfaces.emotion import demo as emotion_demo
from interfaces.ner import demo as ner_demo
from interfaces.ner import download_models as download_spacy_models
from utils import download_hf_models

with gr.Blocks() as demo:
    gr.Markdown(
        f"""
        <style>
        @import 'https://fonts.googleapis.com/css?family=Source+Sans+Pro:300,400';
        </style>
        <div style="display: block; text-align: left; padding:0; margin:0;font-family: "Source Sans Pro", Helvetica, sans-serif;">
            <h1 style="text-align: center;font-size: 17pt;">Babel Machine Demo</h1>
            <p style="font-size: 14pt;">This is a demo for text classification using language models finetuned on data labeled by <a href="https://www.comparativeagendas.net/">CAP</a>, <a href="https://manifesto-project.wzb.eu/">Manifesto Project</a>, sentiment, emotion coding and Named Entity Recognition systems.
            For the coding of complete datasets, please visit the official <a href="https://babel.poltextlab.com/">Babel Machine</a> site.<br>
            Please note that the sentiment (3) and emotions (6) models have been trained using parliamentary speech data, so the results for generic sentences may not be reliable. Additionally, be aware that named entity inputs are case sensitive.<br>
            <br>
            The models listed for Manifesto, Sentiment (3) and Emotions (6) tasks are a beta version and thus not publicly available, 
            the Hugging Face link will not work for them for the time being. We expect a public version after tests and improvements in the Fall. 
            Please feel free to check back for model updates, or reach out to us at that point if you wish to ask about a specific model.
            </p>
        </div>
        """)

    gr.TabbedInterface(
        interface_list=[cap_demo, manifesto_demo, sentiment_demo, emotion_demo, ner_demo],
        tab_names=["CAP", "Manifesto", "Sentiment (3)", "Emotions (6)", "Named Entity Recognition"],
    )

if __name__ == "__main__":
    download_hf_models()
    download_spacy_models()
    demo.launch()

# TODO: add all languages & domains