model bug fix

app.py

@@ -28,10 +28,11 @@ def get_models(llama=False):
     access_token = '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'
     llama_tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=access_token, use_fast=True)#, use_fast=True)
     llama_model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=access_token, device_map={'':0})#, load_in_4bit=True)
-
-  st.write("The assistant is loaded and ready to use!")
-
-  return model, tokenizer, llama_model, llama_tokenizer if llama else model, tokenizer
+    st.write("The assistant is loaded and ready to use!")
+    return model, tokenizer, llama_model, llama_tokenizer 
+  
+  else:
+    return model, tokenizer
 
 model, tokenizer = get_models()