Update app.py

app.py

@@ -8,25 +8,24 @@ import gradio as gr
 # checkpoint = "tiiuae/falcon-7b"
 # checkpoint = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 # checkpoint = "gpt2"
-checkpoint = "vsrinivas/FalconLite2"
+checkpoint = "amazon/FalconLite2"
 
 model = AutoModelForCausalLM.from_pretrained(
     checkpoint, device_map="auto", 
     offload_folder="off_load", 
-    # trust_remote_code=True,
+    trust_remote_code=True,
     # torch_dtype="auto",
 )
-# tokenizer = AutoTokenizer.from_pretrained(checkpoint, 
-#                                           trust_remote_code=True,
-#                                           torch_dtype="auto"
-#                                          )
-
-# model = "tiiuae/FalconLite2"
-tokenizer = AutoTokenizer.from_pretrained(model, 
-                                          # trust_remote_code=True,
+tokenizer = AutoTokenizer.from_pretrained(checkpoint, 
+                                          trust_remote_code=True,
                                           torch_dtype="auto"
                                          )
 
+# model = "tiiuae/FalconLite2"
+# tokenizer = AutoTokenizer.from_pretrained(model, 
+#                                           torch_dtype="auto"
+#                                          )
+
 pipeline = transformers.pipeline(
     "text-generation",
     model=model,