Update app.py

app.py

@@ -6,19 +6,10 @@ import spaces
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from bitsandbytes.functional import quantize_blockwise
-from bitsandbytes.nn import Linear4bit
 
 tokenizer = AutoTokenizer.from_pretrained("FlawedLLM/Bhashini_gemma_merged4bit_clean_final")
-model = AutoModelForCausalLM.from_pretrained("FlawedLLM/Bhashini_gemma_merged4bit_clean_final")
-def quantize_model(model):
-    for name, module in model.named_modules():
-        # Quantize only Linear layers that haven't already been quantized
-        if isinstance(module, torch.nn.Linear) and not isinstance(module, Linear4bit):
-            module = quantize_blockwise(module)
+model = AutoModelForCausalLM.from_pretrained("FlawedLLM/Bhashini_gemma_merged4bit_clean_final", load_in_4bit=True, device_map="auto")
 
-# Quantize the model (modified)
-quantize_model(model)
 # alpaca_prompt = You MUST copy from above!
 @spaces.GPU(duration=300)
 def chunk_it(input_command, item_list):