import streamlit as st
from transformers import pipeline

# 加载模型
@st.cache_resource
def load_pipelines():
    sentiment_pipeline = pipeline(model="KeviniveK/CustomModel_IMDB")
    translation_pipeline = pipeline(
        "translation_en_to_zh",
        model="Helsinki-NLP/opus-mt-en-zh"
    )
    return sentiment_pipeline, translation_pipeline

sentiment_pipeline, translation_pipeline = load_pipelines()

# 标签映射
label_map = {
    "LABEL_0": "差评 | Negative",
    "LABEL_1": "好评 | Positive"
}

# 页面标题
st.title("影评分析与翻译 | Sentiment Analysis & Translation")
st.write("请输入一段英文影评，我们将分析其情感并翻译成中文。")
st.write("Enter an English movie review below. The app will analyze its sentiment and translate it into Chinese.")

# 用户输入
user_input = st.text_area("英文影评输入 | Enter English Movie Review", height=150)

if user_input:
    # 情感分析
    result = sentiment_pipeline(user_input)
    sentiment_raw = result[0]["label"]
    sentiment = label_map.get(sentiment_raw, sentiment_raw)
    confidence = result[0]["score"]

    # 翻译
    translation = translation_pipeline(user_input)
    translated_text = translation[0]["translation_text"]

    # 显示结果
    st.subheader("情感分析结果 | Sentiment Analysis Result")
    st.write(f"**情感 (Sentiment):** {sentiment}")
    st.write(f"**置信度 (Confidence):** {confidence:.2f}")

    st.subheader("中文翻译结果 | Chinese Translation")
    st.write(translated_text)

if __name__ == "__main__":
    pass