Choosing the 10b model loaded in 8 bits

app.py

@@ -12,12 +12,15 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import gradio as gr
 
 #
-# Load the "small" MADLAD400 model (3 billion parameters)
+# Load the "small" MADLAD400 model
 #
-model_name = "google/madlad400-3b-mt"
+model_name = "google/madlad400-10b-mt"
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
 model = AutoModelForSeq2SeqLM.from_pretrained(
-    model_name, device_map="auto", torch_dtype=torch.float16)
+    model_name,
+    device_map="auto",
+    torch_dtype=torch.float16,
+    load_in_8bit=True)
 model = torch.compile(model)
 
 #

requirements.txt

@@ -2,6 +2,7 @@
 torch
 transformers
 accelerate
+bitsandbytes
 sentencepiece
 tokenizers
 optimum