Update app.py

app.py

@@ -1,9 +1,14 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
+# from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import LlamaTokenizer, LlamaForCausalLM
 
 # Load your model here
-tokenizer = AutoTokenizer.from_pretrained("qbwmwsap/amber-model-mine")
-model = AutoModelForCausalLM.from_pretrained("qbwmwsap/amber-model-mine")
+# tokenizer = AutoTokenizer.from_pretrained("qbwmwsap/amber-model-mine")
+# model = AutoModelForCausalLM.from_pretrained("qbwmwsap/amber-model-mine")
+
+# Load original amber model
+tokenizer = LlamaTokenizer.from_pretrained("LLM360/Amber", revision="ckpt_356")
+model = LlamaForCausalLM.from_pretrained("LLM360/Amber", revision="ckpt_356")
 
 st.title('Amber-Model-Mine')