quantization

app/embedding_setup.py

@@ -2,7 +2,7 @@ from langchain_community.vectorstores import Chroma
 from langchain_community.embeddings import HuggingFaceBgeEmbeddings
 
 from langchain.docstore.document import Document
-from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig,BitsAndBytesConfig
 from peft import PeftModel
 from config import *
 import os
@@ -34,9 +34,13 @@ retriever = db.as_retriever(search_kwargs={"k": TOP_K})
 lora_weights_rec = REC_LORA_MODEL
 lora_weights_exp = EXP_LORA_MODEL
 hf_auth  = os.environ.get("hf_token")
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_quant_type="nf4"
+)
 
-
-tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL, token=hf_auth)
+tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL, quantization_config=quantization_config, token=hf_auth)
 
 
 first_token = 'First'
@@ -130,7 +134,7 @@ def compare_docs_with_context(doc_a, doc_b, target_occupation_name, target_occup
     
     #courses = f"First: name: {doc_a.metadata['name']}  description:{doc_a.metadata['description']} Second: name: {doc_b.metadata['name']}  description:{Sdoc_b.metadata['description']}" 
     courses = f"First: name: {doc_a.metadata['name']}  learning outcomes:{doc_a.metadata['skills'][:2000]} Second: name: {doc_b.metadata['name']}  learning outcomes:{doc_b.metadata['skills'][:2000]}" 
-    target_occupation = f"name: {target_occupation_name} description: {target_occupation_dsp}"
+    target_occupation = f"name: {target_occupation_name} description: {target_occupation_dsp[:2000]}"
     skill_gap = skill_gap
     prompt = generate_prompt(target_occupation, skill_gap, courses)
     prompt = [prompt]