Adding 4-bit support

app.py

@@ -4,7 +4,7 @@ import os
 import streamlit as st
 import torch
 from dotenv import load_dotenv
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 from hangman import guess_letter
 from hf_utils import query_hint, query_word
@@ -34,6 +34,8 @@ def setup(model_id: str, device: str) -> None:
         model_id (str): Model ID used to load the tokenizer and model.
     """
     logger.info(f"Loading model and tokenizer from model: '{model_id}'")
+    quantization_config = BitsAndBytesConfig(load_in_4bit=True)
+
     tokenizer = AutoTokenizer.from_pretrained(
         model_id,
         token=os.environ["HF_ACCESS_TOKEN"],
@@ -42,6 +44,7 @@ def setup(model_id: str, device: str) -> None:
         model_id,
         torch_dtype=torch.float16,
         token=os.environ["HF_ACCESS_TOKEN"],
+        quantization_config=quantization_config,
     ).to(device)
     logger.info("Setup finished")
     return {"tokenizer": tokenizer, "model": model}