Update app.py

app.py

@@ -359,14 +359,11 @@ def extract_answer(inputs, outputs, tokenizer):
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 INPUT_MAX_LEN = 128  # Adjusted input length
-OUTPUT_MAX_LEN = 512  # Adjusted output length
+OUTPUT_MAX_LEN = 256  # Adjusted output length
 
-@st.cache_data
-def download_model_name():
-    MODEL_NAME = "VietAI/vit5-base"
-    return MODEL_NAME
+MODEL_NAME = "VietAI/vit5-base"
+ 
 
-MODEL_NAME = download_model_name()
 
 tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME, model_max_length=INPUT_MAX_LEN)
 
@@ -403,16 +400,14 @@ class T5Model(pl.LightningModule):
         return AdamW(self.parameters(), lr=0.0001)
 
 
-@st.cache_data
-def load_t5():
-    train_model = T5Model.load_from_checkpoint('./data-law/law-model-v1.ckpt')
-    train_model.freeze()
-    return train_model
 
-train_model = load_t5()
+train_model = T5Model.load_from_checkpoint('./data-law/law-model-v1.ckpt')
+train_model.freeze()
+
 
 
 def generate_question(question):
+    print("tokenizer")
     inputs_encoding = tokenizer(
         question,
         add_special_tokens=True,
@@ -423,6 +418,7 @@ def generate_question(question):
         return_tensors="pt"
     ).to(DEVICE)
 
+    print("generate id")
     generate_ids = train_model.model.generate(
         input_ids=inputs_encoding["input_ids"],
         attention_mask=inputs_encoding["attention_mask"],
@@ -432,7 +428,8 @@ def generate_question(question):
         no_repeat_ngram_size=2,
         early_stopping=True,
     )
-
+    
+    print("decode")
     preds = [
         tokenizer.decode(gen_id, skip_special_tokens=True, clean_up_tokenization_spaces=True)
         for gen_id in generate_ids