Update app.py

app.py

@@ -361,7 +361,11 @@ DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 INPUT_MAX_LEN = 128  # Adjusted input length
 OUTPUT_MAX_LEN = 256  # Adjusted output length
 
-MODEL_NAME = "VietAI/vit5-base"
+@st.cache_data
+def download_model_name():
+    MODEL_NAME = "VietAI/vit5-base"
+
+download_model_name()
 
 tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME, model_max_length=INPUT_MAX_LEN)
 
@@ -397,8 +401,13 @@ class T5Model(pl.LightningModule):
     def configure_optimizers(self):
         return AdamW(self.parameters(), lr=0.0001)
 
-train_model = T5Model.load_from_checkpoint('./data-law/law-model-v1.ckpt')
-train_model.freeze()
+
+@st.cache_data
+def load_t5():
+    train_model = T5Model.load_from_checkpoint('./data-law/law-model-v1.ckpt')
+    train_model.freeze()
+
+load_t5()
 
 
 def generate_question(question):
@@ -428,6 +437,7 @@ def generate_question(question):
     ]
 
     response = " ".join(preds[0].split())
+    print('T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5T5')
     return response
 
 # st.title("Chatbot Roberta")