Update src/streamlit_app.py

src/streamlit_app.py

@@ -10,7 +10,7 @@ from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain_community.llms import Ollama
 import os
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
 from huggingface_hub import InferenceClient
 import re
 
@@ -67,11 +67,11 @@ def load_retriever():
 def load_llm():
     # pipe = pipeline("text-generation", model="google/flan-t5-small", max_new_tokens=256)
     # load the tokenizer and model on cpu/gpu
-
+    quantization_config = BitsAndBytesConfig(load_in_8bit=True,llm_int8_enable_fp32_cpu_offload=True)
     model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
     # model_name = "meta-llama/Llama-2-7b-chat-hf"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", device_map="auto", load_in_8bit=True)
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", quantization_config=quantization_config, device_map="auto", low_cpu_mem_usage=True)
     pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=256)
     
     return HuggingFacePipeline(pipeline=pipe)