Quantization in Model

app.py

@@ -2,7 +2,7 @@ import os
 import streamlit as st
 import pickle
 import time
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from langchain.llms.base import LLM
 from langchain.chains import RetrievalQAWithSourcesChain
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -16,7 +16,14 @@ login(os.getenv('HF_llama3chat8b'))
 
 class CustomHuggingFaceLLM(LLM):
     def __init__(self, model_name, temperature=0.7):
-        self.model = AutoModelForCausalLM.from_pretrained(model_name)
+
+        # Configure 8-bit quantization using `BitsAndBytesConfig`
+        quantization_config = BitsAndBytesConfig(
+            load_in_8bit=True,  # Enable 8-bit quantization
+            llm_int8_enable_fp32_cpu_offload=True  # Offload FP32 operations to CPU for further memory savings
+        )
+        
+        self.model = AutoModelForCausalLM.from_pretrained(model_name,quantization_config=quantization_config)
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.temperature = temperature