log

app.py

@@ -1,12 +1,34 @@
 import spaces
-from transformers import pipeline
+from transformers import pipeline, BitsAndBytesConfig, AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
+import torch
+import logging
 
-pipe = pipeline(model="ejschwartz/decaf-v1-22b-4bit")
-pipe.model.to("cuda")
+logging.basicConfig(level=logging.INFO)
+logging.info(f"CUDA available: {torch.cuda.is_available()}, CUDA version: {torch.version.cuda}")
+try:
+    import bitsandbytes as _bnb
+    logging.info(f"bitsandbytes version: {_bnb.__version__}")
+except Exception as e:
+    logging.warning(f"Could not import bitsandbytes: {e}")
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,  # key
+)
+
+tokenizer = AutoTokenizer.from_pretrained("ejschwartz/decaf-v1-22b-4bit")
+model = AutoModelForCausalLM.from_pretrained(
+    "ejschwartz/decaf-v1-22b-4bit",
+    device_map="auto",
+    quantization_config=bnb_config,
+)
+
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 @spaces.GPU(size="xlarge")
 def generate(text):
+    print(f"Generating text... {text}")
     return pipe(text)[0]
 
 demo = gr.Interface(fn=generate, inputs="text", outputs="text")