Tasya/laba2.2

1import streamlit as st2 3from transformers import AutoTokenizer, DistilBertForSequenceClassification4import torch5from torch.nn.functional import softmax6 7 8@st.cache9def load_tags_info():10    tag_id = {}11    id_tag = {}12    with open('tag.txt', 'r') as file:13        i = 014        for line in file:15            tag = line[1:].split(',')[0]16            17            tag_id[tag] = i18            id_tag[i] = tag19            i += 120    21    tag_id['None'] = 15022    id_tag[150] = 'None'23 24    return (tag_id, id_tag)25 26tag_id, id_tag = load_tags_info()27 28@st.cache29def load_tokenizer():30    return AutoTokenizer.from_pretrained('./')31 32@st.cache33def load_model():34    return DistilBertForSequenceClassification.from_pretrained('./')35 36def top_xx(preds, xx=95):37    tops = torch.argsort(preds, 1, descending=True)38    total = 039    index = 040    result = []41    while total < xx / 100:42        next_id = tops[0, index].item()43        if next_id == 150:44            index += 145            continue46        total += preds[0, next_id]47        index += 148        result.append(id_tag[next_id])49    return result50    51 52model = load_model()53tokenizer = load_tokenizer()54 55title = st.text_area(label='Title', height=30)56abstract = st.text_area(label='Abstract (optional)', height=200)57st.caption('Generation:')58 59prompt = 'Title: ' + title + ' Abstract: ' + abstract60tokens = tokenizer(prompt, truncation=True, padding='max_length', return_tensors='pt')['input_ids']61preds = softmax(model(tokens.reshape(1, -1)).logits , dim=1)62tags = top_xx(preds)63other_tags = []64st.header('Inferred tags:')65for i, tag_data in enumerate(tags):66    st.caption(tag_data)