| 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) | |