import streamlit as st
from transformers import pipeline
import pandas as pd

# 缓存模型加载
@st.cache_resource
def load_sentiment_model():
    return pipeline("text-classification", model="jinchenliuljc/ecommerce-sentiment-analysis")

@st.cache_resource
def load_ner_model():
    return pipeline("ner", model="jinchenliuljc/ecom_ner_model")

# 处理NER结果
def extract_products(ner_result):
    products = []
    current_product = None
    
    for entity in ner_result:
        if entity['entity'] == 'B-HCCX':
            if current_product is not None:
                products.append(current_product)
            current_product = {
                'start': entity['start'],
                'end': entity['end'],
                'text': entity['word']
            }
        elif entity['entity'] == 'I-HCCX' and current_product is not None:
            # 合并连续字符（中文按字处理）
            current_product['end'] = entity['end']
            current_product['text'] += entity['word']
    
    if current_product is not None:
        products.append(current_product)
    
    return [p['text'] for p in products]

# 初始化session state
if 'records' not in st.session_state:
    st.session_state.records = []

# 页面布局
st.title("DTC客户评论分析系统")
col1, col2 = st.columns(2)

with col1:
    user_input = st.text_input("请输入客户评论：", key="input")
    
    if user_input:
        # 情感分析
        sentiment_classifier = load_sentiment_model()
        sentiment_result = sentiment_classifier(user_input)[0]['label']
        
        # 处理结果
        if sentiment_result == 'LABEL_0':
            st.success("感谢您的积极反馈！❤️")
        else:
            # NER分析
            ner_pipe = load_ner_model()
            ner_result = ner_pipe(user_input)
            
            # 提取产品
            products = extract_products(ner_result)
            
            if products:
                # 添加到记录
                for product in products:
                    new_record = {
                        '产品类别': product,
                        '评论内容': user_input
                    }
                    st.session_state.records.append(new_record)
                
                st.warning(f"检测到问题产品：{', '.join(products)}")
            else:
                st.warning("未识别到具体产品")

with col2:
    if st.session_state.records:
        df = pd.DataFrame(st.session_state.records)
        st.dataframe(
            df,
            column_config={
                "产品类别": "问题产品",
                "评论内容": "相关评论"
            },
            hide_index=True,
            use_container_width=True
        )
    else:
        st.info("暂无客户反馈记录")