import streamlit as st
from transformers import pipeline
import launchdarkly_api

# Initialize LaunchDarkly client
ld_client = launchdarkly_api.LDClient("YOUR_LAUNCHDARKLY_SDK_KEY")

# Model descriptions
model_descriptions = {
    "bert-base-uncased": "BERT base model (uncased)",
    "roberta-base": "RoBERTa base model",
    "distilbert-base-uncased": "DistilBERT base model (uncased)",
    "albert-base-v2": "ALBERT base model v2"
}

# Create a function to get the active model from LaunchDarkly
def get_active_model():
    if ld_client.variation("use_bert", {"key": "user"}):
        return pipeline("sentiment-analysis", model="bert-base-uncased"), "bert-base-uncased"
    elif ld_client.variation("use_roberta", {"key": "user"}):
        return pipeline("sentiment-analysis", model="roberta-base"), "roberta-base"
    elif ld_client.variation("use_distilbert", {"key": "user"}):
        return pipeline("sentiment-analysis", model="distilbert-base-uncased"), "distilbert-base-uncased"
    elif ld_client.variation("use_albert", {"key": "user"}):
        return pipeline("sentiment-analysis", model="albert-base-v2"), "albert-base-v2"
    else:
        return pipeline("sentiment-analysis", model="distilbert-base-uncased"), "distilbert-base-uncased"  # Default model

# Streamlit app
st.title("Sentiment Analysis Demo")
user_input = st.text_area("Enter text for sentiment analysis:")

if st.button("Analyze"):
    model, model_name = get_active_model()
    result = model(user_input)
    st.write(f"Model used: {model_descriptions[model_name]}")
    st.write(result)