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": "Sandaruth/T5_Full_Fine_Tuned_FINDSUM",
    "BERT": "bert-base-uncased",  # Note: BERT isn't designed for summarization; you can change this
    "LongT5": "google/long-t5-local-base",
    "Pegasus": "google/pegasus-xsum"
}

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

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

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

# Button to generate the summary
if st.button("Generate Summary"):
    # Load the selected model and summarizer pipeline
    summarizer = pipeline("summarization", model=models[model_choice])
    
    if input_text:
        # Generate the summary
        summary = summarizer(input_text, max_length=350, min_length=30, do_sample=False)
        
        # Display the summary
        st.subheader("Generated Summary")
        st.write(summary[0]['summary_text'])
    else:
        st.write("Please enter text to summarize!")