Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM, pipeline, logging
+from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM, pipeline, logging, AutoModelForCausalLM
 import torch
 import os
 import httpx
@@ -136,7 +136,9 @@ if submit_button:
         except Exception as error:
             translated_text = error
     if model_name == "winninghealth/WiNGPT-Babel-2":
-        pipe = pipeline("translation", model="winninghealth/WiNGPT-Babel-2")
+        tokenizer = AutoTokenizer.from_pretrained("winninghealth/WiNGPT-Babel-2")
+        model = AutoModelForCausalLM.from_pretrained("winninghealth/WiNGPT-Babel-2")
+        pipe = pipeline("translation", model=model, tokenizer=tokenizer)
         translation = pipe(input_text)
         translated_text = translation[0]['translation_text']
     # Display the translated text