askdmad

handler.py

@@ -3,16 +3,26 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import subprocess
 import importlib
 
-def install_and_import(package):
+import subprocess
+import importlib
+
+def install_package(package):
+    subprocess.check_call(["pip", "install", package])
+
+def install_and_import(package, index_url=None):
     try:
         importlib.import_module(package)
     except ImportError:
-        subprocess.check_call(["pip", "install", package])
+        if index_url:
+            subprocess.check_call(["pip", "install", package, "--index-url", index_url])
+        else:
+            subprocess.check_call(["pip", "install", package])
     finally:
         globals()[package] = importlib.import_module(package)
 
-# Ensure 'bitsandbytes' is installed
-install_and_import('bitsandbytes')
+# Install packages
+install_package('xformers --index-url https://download.pytorch.org/whl/cu121')
+install_package('unsloth[kaggle-new] @ git+https://github.com/unslothai/unsloth.git')
 class EndpointHandler():
     def __init__(self, path=""):
       model_id = "DisgustingOzil/Academic-MCQ-Generator"