Update app.py

app.py

@@ -4,24 +4,29 @@ import torch
 import gradio as gr
 
 desired_dtype = torch.bfloat16
+torch.set_default_dtype(torch.bfloat16)
 
-# checkpoint  = "tiiuae/falcon-40b-instruct"
-checkpoint ="tiiuae/falcon-7b-instruct"
-# checkpoint =  "tiiuae/falcon-7b"
-# checkpoint = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-# checkpoint = "gpt2"
-# checkpoint = "amazon/FalconLite2"
+
+# checkpoint = "vsrinivas/falconlite2"
+checkpoint = "tiiuae/falcon-7b-instruct"
 
 model = AutoModelForCausalLM.from_pretrained(
-    checkpoint, device_map="auto", 
-    offload_folder="off_load", 
-    trust_remote_code=True,
-    # torch_dtype="auto",
-)
-tokenizer = AutoTokenizer.from_pretrained(checkpoint, 
-                                          trust_remote_code=True,
-                                          torch_dtype="auto",
-                                         )
+    # checkpoint, device_map="auto", offload_folder="offload", trust_remote_code=True, torch_dtype="auto")
+    checkpoint, device_map="auto", offload_folder="offload", trust_remote_code=True)
+
+# tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True, torch_dtype="auto")
+tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
+
+# model = AutoModelForCausalLM.from_pretrained(
+#     checkpoint, device_map="auto", 
+#     # offload_folder="off_load", 
+#     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, 
@@ -32,14 +37,23 @@ pipeline = transformers.pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    # use_safetensors=True, 
-    # torch_dtype=torch.bfloat16,
+    torch_dtype=torch.bfloat16,
     trust_remote_code=True,
     device_map="auto",
-    offload_folder="off_load",
-    # offload_state_dict = True,
 )
 
+# pipeline = transformers.pipeline(
+#     "text-generation",
+#     model=model,
+#     tokenizer=tokenizer,
+#     # use_safetensors=True, 
+#     # torch_dtype=torch.bfloat16,
+#     trust_remote_code=True,
+#     device_map="auto",
+#     offload_folder="off_load",
+#     # offload_state_dict = True,
+# )
+
 # def format_chat_prompt(message, chat_history):
 #     prompt = ""
 #     for turn in chat_history: