Update agent.py

agent.py

@@ -54,7 +54,7 @@ from langchain.chains import LLMChain
 from langchain.agents import initialize_agent, Tool, AgentType
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from huggingface_hub import login
-
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
 
 load_dotenv()
 
@@ -461,16 +461,33 @@ tools = [wiki_tool, calc_tool, file_tool, web_tool, arvix_tool, youtube_tool, vi
 
 
 # Get the Hugging Face API token from the environment variable
+# Load Hugging Face token
 hf_token = os.getenv("HF_TOKEN")
-
 login(token=hf_token)
 
-# Initialize the desired model and parameters
+# Model name
 model_name = "mistralai/Mistral-7B-Instruct-v0.1"
+
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Create a text generation pipeline
+# 🔧 Use 4-bit or 8-bit quantization
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,             # or use load_in_8bit=True
+    bnb_4bit_compute_dtype="float16",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+)
+
+# Load model with quantization
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+    quantization_config=bnb_config,
+    token=hf_token
+)
+
+# Create pipeline
 pipe = pipeline(
     "text-generation",
     model=model,
@@ -481,12 +498,10 @@ pipe = pipeline(
     repetition_penalty=1.15
 )
 
-# Create LangChain LLM wrapper
+# LangChain wrapper
 llm = HuggingFacePipeline(pipeline=pipe)
 
-
-
-# Initialize the LangChain agent with the tool(s) and the model
+# Initialize LangChain agent
 agent = initialize_agent(
     tools=tools,
     llm=llm,
@@ -496,6 +511,7 @@ agent = initialize_agent(
 
 
 
+
 # -------------------------------
 # Step 8: Use the Planner, Classifier, and Decision Logic
 # -------------------------------