Create app.py

app.py

@@ -0,0 +1,37 @@
+from datasets import Dataset
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, Seq2SeqTrainer, TrainingArguments
+import streamlit as st
+
+@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)