import streamlit as st
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer


def combine_title_summary(title, summary):
    return "title: " + title + " summary: " + summary


tag2ind = {
    "bio": 0,
    "physics": 1,
    "math": 2,
    "cs": 3,
}


# @st.cache_resource
# def load_model():
#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#     # assert torch.cuda.is_available()
#     tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-cased")
#     model = AutoModelForSequenceClassification.from_pretrained(
#         "./my_model/checkpoint-513"
#     ).to(device)

#     return tokenizer, model


# tokenizer, model = load_model()


# def run_model(model, tokenizer, title, summary):
#     text = combine_title_summary(title, summary)

#     tokens_info = tokenizer(
#         text,
#         padding=False,
#         truncation=True,
#         return_tensors="pt",
#     )

#     model.eval()
#     model.cpu()
#     with torch.no_grad():
#         out = model(**tokens_info)
#         probs = torch.nn.functional.softmax(out.logits, dim=-1)[0]

#         result = f"Text: `{text}`\nPrediction (prob): \n" + "\n".join(
#             [f"{tag}={tag_prob}" for tag, tag_prob in zip(tag2ind, probs)]
#         )
#         return result


title = st.text_input(label="Title", value="")
abstract = st.text_input(label="Abstract", value="")
if st.button("Submit"):
    if title == "" and abstract == "":
        st.error("At least one of title or abstract must be provided")
    else:
        result = combine_title_summary(title, abstract)
        st.success(result)

        # result = run_model(model, tokenizer, title, abstract)
        # st.success(result)