Update app.py

app.py

@@ -11,7 +11,11 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # Summarization model
 summarizer_model_id = "facebook/bart-large-cnn"
 summarizer_tokenizer = SummarizerTokenizer.from_pretrained(summarizer_model_id)
-summarizer_model = AutoModelForSeq2SeqLM.from_pretrained(summarizer_model_id)
+summarizer_model = AutoModelForSeq2SeqLM.from_pretrained(
+    summarizer_model_id,
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
 summarizer_model.to(device)
 
 def summarize_text(text: str, max_length=150) -> str:
@@ -33,6 +37,7 @@ model_id = "MindVR/JohnTran_Fine-tune"
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
+    torch_dtype=torch.float16,
     device_map="auto",
     low_cpu_mem_usage=True,
     token=HF_TOKEN
@@ -61,7 +66,7 @@ def chat(
     with torch.no_grad():
         output = model.generate(
             input_ids,
-            max_new_tokens=1000,
+            max_new_tokens=256,
             do_sample=True,
             top_p=0.95,
             temperature=0.7,