Update main.py

main.py

@@ -1,4 +1,4 @@
-import os
+
 import tempfile
 import requests
 from fastapi import FastAPI, HTTPException, Header, Request
@@ -14,10 +14,10 @@ from sentence_transformers import SentenceTransformer
 import faiss
 import numpy as np
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+# Load environment variables
 
 import os
 os.environ["HF_HOME"] = "./cache"
-# Load environment variables
 load_dotenv()
 HF_API_TOKEN = os.getenv("HF_API_TOKEN")
 API_KEY = os.getenv("API_KEY")
@@ -37,7 +37,12 @@ EMBED_MODEL = SentenceTransformer("all-MiniLM-L6-v2")
 model_name = "deepseek-ai/deepseek-llm-7b-base"
 hf_token = os.getenv("HF_API_TOKEN")  # Make sure your .env has HF_API_TOKEN
 tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", token=hf_token)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+    token=hf_token,
+    offload_folder="./cache/offload"
+)
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 def query_llm(question: str, context_chunks: list):