import streamlit as st
from transformers import AutoModel, AutoTokenizer
from peft import PeftModel

# Loading LED IN Model
base_model = "unsloth/llama-3-8b-bnb-4bit"
led = AutoModel.from_pretrained(base_model)
adapter_model_in = f"sloganLLama"
led_in = PeftModel.from_pretrained(led, adapter_model_in)
led_in_tokenizer = AutoTokenizer.from_pretrained(base_model)

# Generating Summary
def summarize(model, tokenizer, text):
    input_tokenized = tokenizer.encode(text, return_tensors='pt', max_length=8192, truncation=True)
    summary_ids = model.generate(input_tokenized, num_beams=4, length_penalty=0.1, min_length=32, max_length=512)
    summary = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) for g in summary_ids][0]
    return summary

# Reading Txt File
def read_txt_file(file):
    text = file.read().decode('utf-8')
    return text

st.set_page_config(page_title="Slogan Generation", page_icon="img.png")
title = "Slogan Generation"
col1, col2 = st.columns([1,7])
with col1:
    st.image("img.png")
with col2: st.title(title)
st.write("Capturing attention, conveying value, and driving brand loyalty through impactful slogan.")

if "user_text" not in st.session_state:
    st.session_state.user_text = ""

upload_file = st.file_uploader("Upload a .txt file", type="txt")

if upload_file is not None:
    user_text = read_txt_file(upload_file)
else:
    user_text = st.text_area("Paste your brand description here:", value=st.session_state.user_text, height=300)

if st.button("Generate Slogan"):
    with st.spinner("Generating Slogan..."):
        try:
            summary_text = summarize(led_in, led_in_tokenizer, user_text)
            st.session_state.user_text = user_text
            st.write("")
            st.success(summary_text)
            print(summary_text)
        except Exception as e:
            st.error(f"An error occurred: {e}")