Update app.py

app.py

@@ -31,9 +31,12 @@ MAX_LENGTH = 256
 
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 base_model = AutoModelForSeq2SeqLM.from_pretrained(
     BASE_MODEL,
-    quantization_config=BitsAndBytesConfig(load_in_8bit=True),
+    quantization_config=BitsAndBytesConfig(load_in_8bit=True) if device == "cuda" else None,
+    device_map={"": 0} if device == "cuda" else None,
     low_cpu_mem_usage=True,
 )
 
@@ -91,7 +94,7 @@ def translate_cached(text, source_lang, target_lang):
 
     adapter_name, tgt_code = config
     start = time.time()
-    device = next(model.parameters()).device
+    dev = next(model.parameters()).device
     translated_paragraphs = []
 
     for paragraph in text.split("\n"):
@@ -111,7 +114,7 @@ def translate_cached(text, source_lang, target_lang):
                 truncation=True,
                 max_length=MAX_LENGTH,
             )
-            inputs = {k: v.to(device) for k, v in inputs.items()}
+            inputs = {k: v.to(dev) for k, v in inputs.items()}
 
             with adapter_lock:
                 model.set_adapter(adapter_name)