Update app.py

app.py

@@ -9,7 +9,18 @@ torch.set_default_device("cuda")
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 
-model = "deepapaikar/katzbot-phi2"
+# model_name = "deepapaikar/katzbot-phi2"
+
+# model = AutoModelForCausalLM.from_from_pretrained(model_name)
+
+
+# Initialize the model and tokenizer
+model = AutoModelForCausalLM.from_pretrained("deepapaikar/katzbot-phi2",
+                                             torch_dtype=torch.float16, 
+                                             device_map="auto",
+                                             trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("deepapaikar/katzbot-phi2", trust_remote_code=True)
+
 
 # pipeline = transformers.pipeline(
 #     "text-generation",
@@ -17,7 +28,7 @@ model = "deepapaikar/katzbot-phi2"
 #     torch_dtype=torch.float16,
 # )
 
-tokenizer = AutoTokenizer.from_pretrained(model)
+# tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 
 # def predict_answer(question, token=25):
@@ -45,7 +56,7 @@ def predict_answer(question, token=25):
     inputs = {k: v.to(device) for k, v in inputs.items()}
     
     # Use the model directly for inference
-    model.eval()  # Ensure the model is in evaluation mode
+    model.eval() # Ensure the model is in evaluation mode
     model.to(device)  # Ensure the model is on the correct device
     
     # Generate outputs