from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import streamlit as st
import torch

@st.cache_resource
def load_model(cp_path):
    model = AutoModelForSeq2SeqLM.from_pretrained(cp_path)
    return model
    
@st.cache_resource
def load_tokenizer(path):
    tokenizer = AutoTokenizer.from_pretrained(path)

cp_aug = 'minnehwg/finetune-newwiki-summarization-ver-augmented2'
cp_org = 'minnehwg/finetune-newwiki-summarization-ver2'

model_org = load_model(cp_org)
model_aug  = AutoModelForSeq2SeqLM.from_pretrained(cp_aug)
tokenizer = load_tokenizer("VietAI/vit5-base")

def summarize(text, model, tokenizer, num_beams=4, device='cpu'):
    model.eval()
    model.to(device)
    inputs = tokenizer.encode(text, return_tensors="pt", max_length=1024, truncation=True, padding = True).to(device)
    with torch.no_grad():
        summary_ids = model.generate(inputs, max_length=256, num_beams=num_beams)
    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return summary

if text:
    re1 = summarize(model_org, tokenizer, text)
    re2 = summarize(model_aug, tokenizer, text)
    out = {
        'Result from model with original data': re1,
        'Result from model with augmented data': re2
    }
    st.json(out)