import streamlit as st
from transformers import (
    AutoModelForSeq2SeqLM,  # Import the Seq2Seq language model class
    AutoTokenizer,           # Import the tokenizer for language models
    GenerationConfig,       # Import the generation configuration for text generation
    TrainingArguments,      # Import training arguments for fine-tuning models
    Trainer                 # Import the Trainer class for model training
)
import torch

# Import the necessary components for PEFT model
from peft import PeftModel, PeftConfig

# Load the base T5 model for the PEFT model
peft_model_base = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base", torch_dtype=torch.bfloat16)

tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")

# Create the PEFT model by loading the pretrained base model and checkpoint
peft_model = PeftModel.from_pretrained(
    peft_model_base,                           # Base model
    './peft-dialogue-summary-checkpoint-local/',  # Checkpoint directory
    torch_dtype=torch.bfloat16,                # PyTorch data type
    is_trainable=False                         # Specify that the model is not trainable
) 

st.title("Dialogue Summarization using Flan T5 model")
with st.expander("Info about Algorithm and Dataset"):
    st.markdown(""" Open-source [FlanT5 model](https://huggingface.co/google/flan-t5-base) 
                has been finetuned on Open-source [data](https://huggingface.co/datasets/knkarthick/dialogsum) 
                using PEFT (Parameter efficient fine-tuning) for Dialogue Summarization """)
    
dialogue = st.text_area(
    "Dialogue for Summarization",
""" #Person1#: Have you considered upgrading your system?
#Person2#: Yes, but I'm not sure what exactly I would need.
#Person1#: You could consider adding a painting program to your software. It would allow you to make up your own flyers and banners for advertising.
#Person2#: That would be a definite bonus.
#Person1#: You might also want to upgrade your hardware because it is pretty outdated now.
#Person2#: How can we do that?
#Person1#: You'd probably need a faster processor, to begin with. And you also need a more powerful hard disc, more memory and a faster modem. Do you have a CD-ROM drive?
#Person2#: No.
#Person1#: Then you might want to add a CD-ROM drive too, because most new software programs are coming out on Cds.
#Person2#: That sounds great. Thanks. """, height=323)

prompt = f""" Summarize the following conversation.

{dialogue}

Summary: """

if st.button("Summarize", type="primary"):
    # Tokenize the prompt and convert it to PyTorch tensors
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids

    peft_model_outputs = peft_model.generate(input_ids=input_ids, generation_config=GenerationConfig(max_new_tokens=200, num_beams=1))
    peft_model_text_output = tokenizer.decode(peft_model_outputs[0], skip_special_tokens=True)

    st.text(peft_model_text_output)