Upload Hindi embeddings model and all associated files

hindi-rag-system.py

@@ -16,6 +16,12 @@ from langchain.vectorstores import FAISS as LangchainFAISS
 from langchain.docstore.document import Document
 from langchain.embeddings.base import Embeddings
 from typing import List, Dict, Any, Optional, Callable
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import gc
+import warnings
+
+# Ignore specific HuggingFace warnings
+warnings.filterwarnings("ignore", category=UserWarning, message=".*The model doesn't have tied token embeddings.*")
 
 # Tokenizer wrapper class - same as in original code
 class SentencePieceTokenizerWrapper:
@@ -702,6 +708,107 @@ def perform_similarity_search(vector_store, query, k=6):
     print(f"Searching for: {query}")
     return vector_store.similarity_search_with_score(query, k=k)
 
+# Llama model loading function
+def load_llama_model(model_name="unsloth/Llama-3.2-1B-Instruct", device="cuda"):
+    """
+    Load and prepare Llama model for text generation
+    """
+    print(f"Loading LLM: {model_name}")
+    
+    # Check if CUDA is available
+    if device == "cuda" and not torch.cuda.is_available():
+        print("CUDA not available, falling back to CPU")
+        device = "cpu"
+
+    # Quantization config for 4-bit precision to save memory
+    quantization = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.float16,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+    ) if device == "cuda" else None
+
+    # Standard HuggingFace loading
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    if device == "cuda":
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            device_map="auto",
+            quantization_config=quantization
+        )
+    else:
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        model = model.to(device)
+    
+    print("Successfully loaded model")
+    
+    return model, tokenizer
+
+def setup_qa_system(model, tokenizer, vector_store):
+    """
+    Set up a direct QA system using the model and retriever
+    """
+    # Create retriever
+    retriever = vector_store.as_retriever(
+        search_type="similarity",
+        search_kwargs={"k": 3}
+    )
+    
+    # Create a function to generate answers
+    def generate_answer(query):
+        # Retrieve documents
+        try:
+            docs = retriever.invoke(query)
+        except:
+            # Fallback to older method if invoke isn't available
+            docs = retriever.get_relevant_documents(query)
+        
+        # Extract the content
+        context = "\n\n".join([doc.page_content for doc in docs])
+        
+        # Create prompt
+        prompt = f"""
+आपको निम्नलिखित संदर्भ से जानकारी के आधार पर एक प्रश्न का उत्तर देना है। 
+यदि आप उत्तर नहीं जानते हैं, तो बस "मुझे नहीं पता" कहें।
+
+संदर्भ:
+{context}
+
+प्रश्न: {query}
+
+उत्तर:
+"""
+        
+        # Generate text
+        inputs = tokenizer(prompt, return_tensors="pt")
+        
+        # Move to the same device as the model
+        for k, v in inputs.items():
+            if hasattr(v, "to") and callable(v.to):
+                inputs[k] = v.to(model.device)
+        
+        with torch.no_grad():
+            try:
+                outputs = model.generate(
+                    inputs.input_ids,
+                    max_new_tokens=512,
+                    temperature=0.7,
+                    top_p=0.9,
+                    do_sample=True
+                )
+            except Exception as e:
+                return f"Error generating response: {str(e)}"
+        
+        # Decode the generated text
+        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract just the answer part (after the prompt)
+        answer = full_response.split("उत्तर:")[-1].strip()
+        
+        return answer
+    
+    return generate_answer
+
 # Main RAG functions
 def index_text_files(model, tokenizer, data_dir, output_dir, device="cuda", chunk_size=500):
     """
@@ -795,7 +902,7 @@ def query_text_corpus(model, tokenizer, vector_store_path, query, k=6, device="c
         
         processed_results.append((combined_doc, score))
     
-    return processed_results
+    return processed_results, vector_store
 
 def main():
     parser = argparse.ArgumentParser(description="Hindi RAG System with LangChain and FAISS")
@@ -821,27 +928,54 @@ def main():
                         help="Run in interactive mode for querying")
     parser.add_argument("--reindex", action="store_true",
                         help="Force reindexing even if index exists")
+    parser.add_argument("--qa", action="store_true",
+                        help="Use LLM for question answering instead of just retrieval")
+    parser.add_argument("--llm_name", type=str, default="unsloth/Llama-3.2-1B-Instruct",
+                        help="HuggingFace model name for the LLM")
     args = parser.parse_args()
     
-    # Load model and tokenizer
-    model, tokenizer, config = load_model_and_tokenizer(args.model_dir, args.tokenizer_dir)
+    # Load embedding model and tokenizer
+    embed_model, embed_tokenizer, config = load_model_and_tokenizer(args.model_dir, args.tokenizer_dir)
     
-    # Move model to device
-    model = model.to(args.device)
+    # Move embedding model to device
+    embed_model = embed_model.to(args.device)
     
     # Create vector store path
     vector_store_path = os.path.join(args.output_dir, "faiss_index")
     
+    # Load LLM if QA is enabled
+    llm_model = None
+    llm_tokenizer = None
+    qa_generator = None
+    
+    if args.qa:
+        try:
+            # Load LLM
+            llm_model, llm_tokenizer = load_llama_model(args.llm_name, args.device)
+            print("LLM loaded successfully for QA")
+        except Exception as e:
+            print(f"Error loading LLM: {e}")
+            print("Falling back to retrieval-only mode")
+            args.qa = False
+    
     if args.index or args.reindex:
         # Index text files
-        index_text_files(model, tokenizer, args.data_dir, args.output_dir, args.device, args.chunk_size)
+        vector_store, _ = index_text_files(
+            embed_model, embed_tokenizer, args.data_dir, args.output_dir, args.device, args.chunk_size
+        )
         print(f"Indexing complete. Vector store saved to {vector_store_path}")
+        
+        # Set up QA chain if enabled
+        if args.qa and llm_model is not None and llm_tokenizer is not None:
+            qa_generator = setup_qa_system(llm_model, llm_tokenizer, vector_store)
     
     if args.query:
         # Query the corpus
-        results = query_text_corpus(model, tokenizer, vector_store_path, args.query, args.top_k, args.device)
+        results, vector_store = query_text_corpus(
+            embed_model, embed_tokenizer, vector_store_path, args.query, args.top_k, args.device
+        )
         
-        # Print results
+        # Print retrieval results
         print("\nSearch Results:")
         for i, (doc, score) in enumerate(results):
             print(f"\nResult {i+1} (Score: {score:.4f}):")
@@ -850,10 +984,27 @@ def main():
             # Extract and print only relevant sentences
             relevant_text = extract_relevant_sentences(doc.page_content, args.query)
             print(f"Content: {relevant_text}")
+        
+        # If QA is enabled, also answer the question using the LLM
+        if args.qa and llm_model is not None and llm_tokenizer is not None:
+            if qa_generator is None:
+                qa_generator = setup_qa_system(llm_model, llm_tokenizer, vector_store)
+            
+            # Get answer from QA chain
+            print("\nGenerating answer using LLM...")
+            try:
+                answer = qa_generator(args.query)
+                print("\nLLM Answer:")
+                print(answer)
+            except Exception as e:
+                print(f"Error generating answer: {e}")
     
     if args.interactive:
         print("\nInteractive mode. Enter queries (or type 'quit' to exit).")
         
+        # For the first query, load vector store
+        vector_store = None
+        
         while True:
             print("\nEnter query:")
             query = input()
@@ -865,9 +1016,11 @@ def main():
                 break
             
             # Query the corpus
-            results = query_text_corpus(model, tokenizer, vector_store_path, query, args.top_k, args.device)
+            results, vector_store = query_text_corpus(
+                embed_model, embed_tokenizer, vector_store_path, query, args.top_k, args.device
+            )
             
-            # Print results
+            # Print retrieval results
             print("\nSearch Results:")
             for i, (doc, score) in enumerate(results):
                 print(f"\nResult {i+1} (Score: {score:.4f}):")
@@ -876,6 +1029,24 @@ def main():
                 # Extract and print only relevant sentences
                 relevant_text = extract_relevant_sentences(doc.page_content, query)
                 print(f"Content: {relevant_text}")
+            
+            # If QA is enabled, also answer the question using the LLM
+            if args.qa and llm_model is not None and llm_tokenizer is not None:
+                if qa_generator is None:
+                    qa_generator = setup_qa_system(llm_model, llm_tokenizer, vector_store)
+                
+                print("\nGenerating answer using LLM...")
+                try:
+                    answer = qa_generator(query)
+                    print("\nLLM Answer:")
+                    print(answer)
+                except Exception as e:
+                    print(f"Error generating answer: {e}")
+    
+    # Clean up GPU memory
+    if args.device == "cuda":
+        gc.collect()
+        torch.cuda.empty_cache()
 
 if __name__ == "__main__":
     main()

hindi-rag-system.py.amltmp

@@ -16,6 +16,12 @@ from langchain.vectorstores import FAISS as LangchainFAISS
 from langchain.docstore.document import Document
 from langchain.embeddings.base import Embeddings
 from typing import List, Dict, Any, Optional, Callable
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import gc
+import warnings
+
+# Ignore specific HuggingFace warnings
+warnings.filterwarnings("ignore", category=UserWarning, message=".*The model doesn't have tied token embeddings.*")
 
 # Tokenizer wrapper class - same as in original code
 class SentencePieceTokenizerWrapper:
@@ -702,6 +708,107 @@ def perform_similarity_search(vector_store, query, k=6):
     print(f"Searching for: {query}")
     return vector_store.similarity_search_with_score(query, k=k)
 
+# Llama model loading function
+def load_llama_model(model_name="unsloth/Llama-3.2-1B-Instruct", device="cuda"):
+    """
+    Load and prepare Llama model for text generation
+    """
+    print(f"Loading LLM: {model_name}")
+    
+    # Check if CUDA is available
+    if device == "cuda" and not torch.cuda.is_available():
+        print("CUDA not available, falling back to CPU")
+        device = "cpu"
+
+    # Quantization config for 4-bit precision to save memory
+    quantization = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.float16,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+    ) if device == "cuda" else None
+
+    # Standard HuggingFace loading
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    if device == "cuda":
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            device_map="auto",
+            quantization_config=quantization
+        )
+    else:
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        model = model.to(device)
+    
+    print("Successfully loaded model")
+    
+    return model, tokenizer
+
+def setup_qa_system(model, tokenizer, vector_store):
+    """
+    Set up a direct QA system using the model and retriever
+    """
+    # Create retriever
+    retriever = vector_store.as_retriever(
+        search_type="similarity",
+        search_kwargs={"k": 3}
+    )
+    
+    # Create a function to generate answers
+    def generate_answer(query):
+        # Retrieve documents
+        try:
+            docs = retriever.invoke(query)
+        except:
+            # Fallback to older method if invoke isn't available
+            docs = retriever.get_relevant_documents(query)
+        
+        # Extract the content
+        context = "\n\n".join([doc.page_content for doc in docs])
+        
+        # Create prompt
+        prompt = f"""
+आपको निम्नलिखित संदर्भ से जानकारी के आधार पर एक प्रश्न का उत्तर देना है। 
+यदि आप उत्तर नहीं जानते हैं, तो बस "मुझे नहीं पता" कहें।
+
+संदर्भ:
+{context}
+
+प्रश्न: {query}
+
+उत्तर:
+"""
+        
+        # Generate text
+        inputs = tokenizer(prompt, return_tensors="pt")
+        
+        # Move to the same device as the model
+        for k, v in inputs.items():
+            if hasattr(v, "to") and callable(v.to):
+                inputs[k] = v.to(model.device)
+        
+        with torch.no_grad():
+            try:
+                outputs = model.generate(
+                    inputs.input_ids,
+                    max_new_tokens=512,
+                    temperature=0.7,
+                    top_p=0.9,
+                    do_sample=True
+                )
+            except Exception as e:
+                return f"Error generating response: {str(e)}"
+        
+        # Decode the generated text
+        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract just the answer part (after the prompt)
+        answer = full_response.split("उत्तर:")[-1].strip()
+        
+        return answer
+    
+    return generate_answer
+
 # Main RAG functions
 def index_text_files(model, tokenizer, data_dir, output_dir, device="cuda", chunk_size=500):
     """
@@ -795,7 +902,7 @@ def query_text_corpus(model, tokenizer, vector_store_path, query, k=6, device="c
         
         processed_results.append((combined_doc, score))
     
-    return processed_results
+    return processed_results, vector_store
 
 def main():
     parser = argparse.ArgumentParser(description="Hindi RAG System with LangChain and FAISS")
@@ -821,27 +928,54 @@ def main():
                         help="Run in interactive mode for querying")
     parser.add_argument("--reindex", action="store_true",
                         help="Force reindexing even if index exists")
+    parser.add_argument("--qa", action="store_true",
+                        help="Use LLM for question answering instead of just retrieval")
+    parser.add_argument("--llm_name", type=str, default="unsloth/Llama-3.2-1B-Instruct",
+                        help="HuggingFace model name for the LLM")
     args = parser.parse_args()
     
-    # Load model and tokenizer
-    model, tokenizer, config = load_model_and_tokenizer(args.model_dir, args.tokenizer_dir)
+    # Load embedding model and tokenizer
+    embed_model, embed_tokenizer, config = load_model_and_tokenizer(args.model_dir, args.tokenizer_dir)
     
-    # Move model to device
-    model = model.to(args.device)
+    # Move embedding model to device
+    embed_model = embed_model.to(args.device)
     
     # Create vector store path
     vector_store_path = os.path.join(args.output_dir, "faiss_index")
     
+    # Load LLM if QA is enabled
+    llm_model = None
+    llm_tokenizer = None
+    qa_generator = None
+    
+    if args.qa:
+        try:
+            # Load LLM
+            llm_model, llm_tokenizer = load_llama_model(args.llm_name, args.device)
+            print("LLM loaded successfully for QA")
+        except Exception as e:
+            print(f"Error loading LLM: {e}")
+            print("Falling back to retrieval-only mode")
+            args.qa = False
+    
     if args.index or args.reindex:
         # Index text files
-        index_text_files(model, tokenizer, args.data_dir, args.output_dir, args.device, args.chunk_size)
+        vector_store, _ = index_text_files(
+            embed_model, embed_tokenizer, args.data_dir, args.output_dir, args.device, args.chunk_size
+        )
         print(f"Indexing complete. Vector store saved to {vector_store_path}")
+        
+        # Set up QA chain if enabled
+        if args.qa and llm_model is not None and llm_tokenizer is not None:
+            qa_generator = setup_qa_system(llm_model, llm_tokenizer, vector_store)
     
     if args.query:
         # Query the corpus
-        results = query_text_corpus(model, tokenizer, vector_store_path, args.query, args.top_k, args.device)
+        results, vector_store = query_text_corpus(
+            embed_model, embed_tokenizer, vector_store_path, args.query, args.top_k, args.device
+        )
         
-        # Print results
+        # Print retrieval results
         print("\nSearch Results:")
         for i, (doc, score) in enumerate(results):
             print(f"\nResult {i+1} (Score: {score:.4f}):")
@@ -850,10 +984,27 @@ def main():
             # Extract and print only relevant sentences
             relevant_text = extract_relevant_sentences(doc.page_content, args.query)
             print(f"Content: {relevant_text}")
+        
+        # If QA is enabled, also answer the question using the LLM
+        if args.qa and llm_model is not None and llm_tokenizer is not None:
+            if qa_generator is None:
+                qa_generator = setup_qa_system(llm_model, llm_tokenizer, vector_store)
+            
+            # Get answer from QA chain
+            print("\nGenerating answer using LLM...")
+            try:
+                answer = qa_generator(args.query)
+                print("\nLLM Answer:")
+                print(answer)
+            except Exception as e:
+                print(f"Error generating answer: {e}")
     
     if args.interactive:
         print("\nInteractive mode. Enter queries (or type 'quit' to exit).")
         
+        # For the first query, load vector store
+        vector_store = None
+        
         while True:
             print("\nEnter query:")
             query = input()
@@ -865,9 +1016,11 @@ def main():
                 break
             
             # Query the corpus
-            results = query_text_corpus(model, tokenizer, vector_store_path, query, args.top_k, args.device)
+            results, vector_store = query_text_corpus(
+                embed_model, embed_tokenizer, vector_store_path, query, args.top_k, args.device
+            )
             
-            # Print results
+            # Print retrieval results
             print("\nSearch Results:")
             for i, (doc, score) in enumerate(results):
                 print(f"\nResult {i+1} (Score: {score:.4f}):")
@@ -876,6 +1029,24 @@ def main():
                 # Extract and print only relevant sentences
                 relevant_text = extract_relevant_sentences(doc.page_content, query)
                 print(f"Content: {relevant_text}")
+            
+            # If QA is enabled, also answer the question using the LLM
+            if args.qa and llm_model is not None and llm_tokenizer is not None:
+                if qa_generator is None:
+                    qa_generator = setup_qa_system(llm_model, llm_tokenizer, vector_store)
+                
+                print("\nGenerating answer using LLM...")
+                try:
+                    answer = qa_generator(query)
+                    print("\nLLM Answer:")
+                    print(answer)
+                except Exception as e:
+                    print(f"Error generating answer: {e}")
+    
+    # Clean up GPU memory
+    if args.device == "cuda":
+        gc.collect()
+        torch.cuda.empty_cache()
 
 if __name__ == "__main__":
     main()