quant

app.py

@@ -17,6 +17,7 @@ def load_model():
     model = AutoModelForCausalLM.from_pretrained(
         base_model_name,
         device_map='auto',
+        quantization_config=nf4_config,
     )
     model.config.use_cache = False
     model.config.pretraining_tp = 1
@@ -49,6 +50,7 @@ display = "Ask Rap-Mistral Something"
 st.write(display)
 question = st.chat_input("Write a verse in the style of Lupe Fiasco")
 if question:
+    display = "Loading..."
     prompt = build_prompt(question)
     inputs = tokenizer(prompt, return_tensors="pt")
     model_inputs = inputs.to('cuda')