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

# Load pipelines
pipe1 = pipeline("translation", model="DunnBC22/opus-mt-zh-en-Chinese_to_English")
pipe3 = pipeline("text-classification", model="lxyuan/distilbert-base-multilingual-cased-sentiments-student")

# Load model and tokenizer for pipe2
tokenizer = AutoTokenizer.from_pretrained("huimanho/CustomModel_Amazon")
model = AutoModelForSequenceClassification.from_pretrained("huimanho/CustomModel_Amazon")

# Streamlit app
st.title("Chinese Review Analysis - Translation, Rating & Sentiment")

# Input text
chinese_text = st.text_area("Enter Chinese Review:", height=150)

if st.button("Analyze"):
    # Translation
    english_text = pipe1(chinese_text)[0]['translation_text']
    
    # Display translation
    st.subheader("Translated Text")
    st.write(english_text)

    # Rating Prediction
    inputs = tokenizer(english_text, return_tensors="pt")
    outputs = model(**inputs)
    prediction = outputs.logits.argmax(-1).item()

    # Display estimated rating
    st.subheader("Estimated Amazon Rating")
    st.write(f"**Rating:** {prediction + 1} out of 5")     

    # Sentiment Classification
    sentiment = pipe3(chinese_text)[0]['label']

    # Display sentiment
    st.subheader("Sentiment Analysis")
    st.write(f"**Sentiment:** {sentiment}")    
    
    # Additional styling
    st.markdown("---")  # Add a horizontal line for separation