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 text = st.text_area('Nhập tiêu đề vào đây') 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)