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

# 加载模型和 tokenizer
model_name3_p2 = "rajesh426/distilbert-base-uncased_Up_Sampling_Sub_Category_SPEECH_TEXT_DISPLAY_v1"
num_labels = 10  # 根据你的数据集设置标签的数量
model3_p2 = AutoModelForSequenceClassification.from_pretrained(model_name3_p2, num_labels=num_labels, ignore_mismatched_sizes=True)
tokenizer3 = AutoTokenizer.from_pretrained(model_name3_p2)

# Streamlit 应用
st.title("Text Classification with HuggingFace Spaces")
st.write("Enter a sentence to classify:")

user_input = st.text_input("")

if user_input:
    inputs = tokenizer3(user_input, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        outputs = model3_p2(**inputs)
    logits = outputs.logits
    predicted_class_id = torch.argmax(logits, dim=-1).item()

    st.write(f"Predicted class ID: {predicted_class_id}")