import streamlit as st
from transformers import pipeline
from huggingface_hub import login
from dotenv import load_dotenv
import os

# Load the environment variables from the .env file
load_dotenv()

# Retrieve the token from the .env file
huggingface_token = os.getenv("HUGGINGFACE_TOKEN")

# Log in using the retrieved token
login(token=huggingface_token)

# Available models for summarization
models = {
    "T5_Full_FineTune_V0.1_40K": "ARSynopsis/T5_Full_FineTune_V0.1_40K",
    "T5_Full_FineTune_V0.1_80K": "ARSynopsis/T5_Full_FineTune_V0.1_80K",
    "BART_Base_Full_FineTune_V0.1_83K": "ARSynopsis/BART_Base_Full_FineTune_V0.1_83K",
    "LongT5": "google/long-t5-local-base",
    "Pegasus": "google/pegasus-xsum"
}

# Function to count words in the input text
def count_words(text):
    return len(text.split())

# Streamlit app layout
st.title("Summarization with Multiple Models")

# Dropdown to select the model (bolded)
st.markdown("### **Select a model for summarization**")
model_choice = st.selectbox("", models.keys())

# Text area for input (bolded)
st.markdown("### **Enter the long text you want to summarize**")
input_text = st.text_area("", height=300)

# Button to generate the summary
if st.button("Generate Summary"):
    # Show a spinner while generating the summary
    with st.spinner("Generating summary, please wait..."):
        # Load the selected model and summarizer pipeline
        summarizer = pipeline("summarization", model=models[model_choice])
        
        # Log the model choice
        st.write(f"Using model: **{model_choice}** for summarization.")
        
        # Count and log the number of words in the input text
        word_count = count_words(input_text)
        st.write(f"Number of words in input: **{word_count}**")

        if input_text:
            # Generate the summary
            summary = summarizer(input_text, max_length=350, min_length=30, do_sample=False)
            
            # Log the success message
            st.success("Summary generated successfully!")
            
            # Display the summary
            st.subheader("Generated Summary")
            st.write(summary[0]['summary_text'])
        else:
            st.warning("Please enter text to summarize!")

# Optionally, you can add a footer or additional instructions
st.markdown("---")
st.write("Provide a long text and select a model to see the summarization in action!")