Update main.py

main.py

@@ -7,6 +7,10 @@ from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from huggingface_hub import InferenceClient  
 from transformers import AutoTokenizer, AutoModel
 from deep_translator import GoogleTranslator
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+from accelerate import infer_auto_device_map
+
 
 
 # Ensure HF_TOKEN is set  
@@ -31,14 +35,14 @@ llm_client = InferenceClient(
 #     generate_kwargs={"temperature": 0.1},  
 # )  
 # Configure Llama index settings with the new model
-Settings.llm = HuggingFaceInferenceAPI(
-    model_name=repo_id,
-    tokenizer_name=repo_id,  # Use the same tokenizer as the model
-    context_window=3000,
-    token=HF_TOKEN,
-    max_new_tokens=512,
-    generate_kwargs={"temperature": 0.1},
-)
+# Settings.llm = HuggingFaceInferenceAPI(
+#     model_name=repo_id,
+#     tokenizer_name=repo_id,  # Use the same tokenizer as the model
+#     context_window=3000,
+#     token=HF_TOKEN,
+#     max_new_tokens=512,
+#     generate_kwargs={"temperature": 0.1},
+# )
 # Settings.embed_model = HuggingFaceEmbedding(  
 #     model_name="BAAI/bge-small-en-v1.5"  
 # )  
@@ -46,17 +50,35 @@ Settings.llm = HuggingFaceInferenceAPI(
 # Settings.embed_model = HuggingFaceEmbedding(
 #     model_name="xlm-roberta-base"  # XLM-RoBERTa model for multilingual support
 # )
-Settings.embed_model = HuggingFaceEmbedding(
-    model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
-)
+# Settings.embed_model = HuggingFaceEmbedding(
+#     model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
+# )
 
 # # Configure tokenizer and model if required
 # tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
 # model = AutoModel.from_pretrained("xlm-roberta-base")
 # Configure tokenizer and model if required
 tokenizer = AutoTokenizer.from_pretrained(repo_id)  # Use the tokenizer from the new model
-model = AutoModel.from_pretrained(repo_id)  # Load the new model
-
+# model = AutoModel.from_pretrained(repo_id)  # Load the new model
+model = AutoModelForCausalLM.from_pretrained(
+    repo_id,
+    load_in_4bit=True,  # Load in 4-bit quantization
+    torch_dtype=torch.float16,
+    device_map="auto",
+)
+# Configure Llama index settings
+Settings.llm = HuggingFaceInferenceAPI(
+    model_name=repo_id,
+    tokenizer_name=repo_id,  # Use the same tokenizer as the model
+    context_window=2048,  # Reduce context window to save memory
+    token=HF_TOKEN,
+    max_new_tokens=256,  # Reduce max tokens to save memory
+    generate_kwargs={"temperature": 0.1},
+)
+# Use a smaller embedding model
+Settings.embed_model = HuggingFaceEmbedding(
+    model_name="sentence-transformers/all-MiniLM-L6-v2"  # Smaller and faster
+)
 PERSIST_DIR = "db"  
 PDF_DIRECTORY = 'data'