Update src/rag.py

src/rag.py

@@ -1,12 +1,3 @@
-# /Users/divyeshpatel/Desktop/archiveWork/rajvi/nlp/rag.py
-# !pip install llama-cpp-python
-
-# from llama.cpp import Llama
-#
-# llm = Llama.from_pretrained(
-# 	repo_id="rdz-falcon/model",
-# 	filename="unsloth.F16.gguf",
-# )
 
 # !pip install langchain
 # !pip install langchain-community
@@ -75,14 +66,20 @@ def load_emotion_classifier(api_base_url="http://127.0.0.1:1234/v1"):
     """
     print(f"=== CONFIGURING LLM CLIENT FOR API: {api_base_url} ===")
 
+    from llama_cpp import Llama
     
-    llm = ChatOpenAI(
-        openai_api_base=api_base_url,
-        openai_api_key="dummy-key", # Required by LangChain, but not used by LM Studio
-        temperature=0.7,
-        max_tokens=128,
+    llm = Llama.from_pretrained(
+    	repo_id="rdz-falcon/llma_fine-tuned",
+    	filename="unsloth.F16.gguf",
     )
-    return llm
+
+    # llm = ChatOpenAI(
+    #     openai_api_base=api_base_url,
+    #     openai_api_key="dummy-key", # Required by LangChain, but not used by LM Studio
+    #     temperature=0.7,
+    #     max_tokens=128,
+    # )
+    # return llm
 
     # --- The following code was commented out or unreachable in the original notebook ---
     # Example code (replace with appropriate code for your model):
@@ -180,6 +177,390 @@ You are an AAC (Augmentative and Alternative Communication) user (Elliot) engagi
 <|user|>
 The conversation partner asked: "{question}"
 
+Please generate your response as the AAC user, following the instructions above.</s>
+<|assistant|>
+
+    """.strip()
+
+    PROMPT = PromptTemplate(
+        input_variables=["question", "emotion_analysis", "context", "chat_history"],
+        template=template,
+    )
+    print("\n Prompt:", PROMPT)
+    return PROMPT
+
+class AACAssistant:
+    def __init__(self, document_path):
+        print("Initializing AAC Assistant...")
+        print("Loading document retriever...")
+        self.vectorstore = setup_document_retriever(document_path)
+        print("Configuring emotion LLM client...")
+        # Use the new function to get the client for the API
+        self.emotion_llm = load_emotion_classifier() # You can pass a different URL if needed
+        print("Loading generation model...")
+        self.llm = load_generation_model() # This now loads the Ollama model
+        print("Creating prompt templates...")
+        self.prompt = create_prompt_templates()
+        print("Setting up conversation memory...")
+
+        # Set up memory for chat history
+        self.memory = ConversationBufferMemory(
+            memory_key="chat_history",
+            return_messages=True,
+            output_key="answer",
+            # Specify the input key for the memory explicitly
+            input_key="question"
+        )
+
+        # Create retrieval chain (using the main generation LLM)
+        self.chain = ConversationalRetrievalChain.from_llm(
+            llm=self.llm, # Use the main generation model here
+            retriever=self.vectorstore.as_retriever(search_kwargs={'k': 3}),
+            memory=self.memory,
+            combine_docs_chain_kwargs={"prompt": self.prompt},
+            return_source_documents=True,
+            verbose=True
+        )
+
+        print("AAC Assistant initialized and ready!")
+
+    def get_emotion_analysis(self,llm, situation):
+        """
+        Gets emotion analysis from the configured emotion LLM API.
+        """
+        # Define the prompt structure for the emotion analysis model
+        # (Adjust this based on how you prompted your model in LM Studio)
+        text = situation
+        response = llm.create_chat_completion(
+        messages=[{"role": "user", "content": text}],
+        max_tokens=128,       # Max length of the generated response (adjust as needed)
+        temperature=0.7,    # Controls randomness (adjust)
+            # top_p=0.9,          # Optional: Nucleus sampling
+            # top_k=40,           # Optional: Top-k sampling
+        stop=["<|eot_id|>"], # Crucial: Stop generation when the model outputs the end-of-turn token
+            stream=False,        # Set to True to get token-by-token output (like TextStreamer)
+        )
+    
+                # --- 4. Extract and print the response ---
+        if response and 'choices' in response and len(response['choices']) > 0:
+            assistant_message = response['choices'][0]['message']['content']
+            print("\nAssistant Response:")
+            print(assistant_message.strip())
+            print("returning:", assistant_message.strip())
+            return assistant_message.strip()
+        else:
+            print("\nNo response generated or unexpected format.")
+            print("Full response:", response)
+
+        return ""
+
+
+    def process_query(self, user_query):
+        """
+        Process a query from the conversation partner to the AAC user.
+
+        Args:
+            user_query (str): Question asked by the conversation partner
+
+        Returns:
+            str: Generated response for the AAC user to communicate
+        """
+        # Step 1: Get emotion analysis from the LM Studio API via the emotion_llm client
+        print(f"Getting emotion analysis for query: '{user_query}'")
+        emotion_analysis = self.get_emotion_analysis(self.emotion_llm, user_query)
+        print(f"Emotion Analysis Result: {emotion_analysis}")
+
+        # Step 2: Run the RAG + LLM chain (using the main generation model)
+        # The emotion_analysis is now passed into the prompt context
+        print("Running main RAG chain...")
+        # Use invoke instead of the deprecated __call__
+        # Pass inputs as a dictionary matching the chain's expected input keys
+        response = self.chain.invoke(
+            {"question": user_query, "emotion_analysis": emotion_analysis}
+        )
+
+        return response["answer"]
+
+# def run_demo():
+#     # Sample personal experiences document path - replace with your actual file
+#     document_path = "aac_user_experiences.txt"
+
+#     # Create a dummy document if it doesn't exist for demonstration
+#     # if not os.path.exists(document_path):
+#     #     with open(document_path, "w") as f:
+#     #         f.write("""
+#     #         I grew up in Seattle and love the rain.
+#     #         My favorite hobby is playing chess, which I've been doing since I was 7 years old.
+#     #         I have a dog named Max who is a golden retriever.
+#     #         I went to college at University of Washington and studied computer science.
+#     #         I enjoy watching sci-fi movies and Star Trek is my favorite series.
+#     #         I've traveled to Japan twice and love Japanese cuisine.
+#     #         Music helps me relax, especially classical piano pieces.
+#     #         I volunteer at the local animal shelter once a month.
+#     #         """)
+
+#     # Initialize the assistant
+#     assistant = AACAssistant(document_path)
+
+#     # Interactive demo
+#     print("\n===== AAC Communication Assistant Demo =====")
+#     print("(Type 'exit' to end the demo)")
+
+#     while True:
+#         try:
+#             user_input = input("\nConversation partner says: ")
+#             if user_input.lower() == 'exit':
+#                 break
+
+#             response = assistant.process_query(user_input)
+#             print(f"\nAAC user communicates: {response}")
+#         except EOFError: # Handle case where input stream ends unexpectedly
+#             print("\nInput stream closed. Exiting demo.")
+#             break
+#         except KeyboardInterrupt: # Handle Ctrl+C
+#             print("\nDemo interrupted by user. Exiting.")
+#             break
+#         except Exception as e:
+#             print(f"\nAn unexpected error occurred: {e}")
+#             # Optionally add more specific error handling or logging
+#             # Consider whether to break or continue the loop on error
+#             break # Exit on error for safety
+
+# try:
+#     from importlib.metadata import PackageNotFoundError
+# except ImportError:
+#     # Define a fallback for older Python versions
+#     class PackageNotFoundError(Exception):
+#         pass
+
+# # Cell 13: Main Execution Block
+# if __name__ == "__main__":
+#     run_demo()
+
+# # !pip install bitsandbytes -q || echo "bitsandbytes installation failed, will use fp16 precision instead"
+# # pip install -U bitsandbytes
+
+# from llama.cpp import Llama
+#
+# llm = Llama.from_pretrained(
+# 	repo_id="rdz-falcon/model",
+# 	filename="unsloth.F16.gguf",
+# )
+
+# !pip install langchain
+# !pip install langchain-community
+
+# !pip install chromadb
+
+import os
+import torch
+import tempfile
+from langchain.chains import ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
+from langchain_community.document_loaders import TextLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import Chroma
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain_community.llms import HuggingFacePipeline
+from langchain.prompts import PromptTemplate
+from langchain_community.llms import Ollama 
+from langchain_openai import ChatOpenAI
+from transformers import AutoTokenizer, pipeline, AutoModelForCausalLM, BitsAndBytesConfig
+
+def setup_document_retriever(document_path):
+    # Load documents with the AAC user's personal experiences
+    loader = TextLoader(document_path)
+    documents = loader.load()
+
+    # Split documents into chunks
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=200,
+        separators=["\n\n", "\n", " ", ""]
+    )
+    chunks = text_splitter.split_documents(documents)
+
+    # Create embeddings
+    embeddings = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-MiniLM-L6-v2",
+        model_kwargs={'device': 'cuda' if torch.cuda.is_available() else 'cpu'}
+    )
+
+    # Create a persistent directory for the ChromaDB
+    persist_directory = os.path.join(tempfile.gettempdir(), "chroma_db")
+
+    # Create Chroma vector store
+    vectorstore = Chroma.from_documents(
+        documents=chunks,
+        embedding=embeddings,
+        persist_directory=persist_directory
+    )
+
+    # Persist the database to disk
+    vectorstore.persist()
+
+    return vectorstore
+
+def load_emotion_classifier(api_base_url="http://127.0.0.1:1234/v1"):
+    """
+    This function configures and returns a LangChain LLM client
+    to interact with an OpenAI-compatible API endpoint (like LM Studio).
+
+    Args:
+        api_base_url (str): The base URL of the OpenAI-compatible API endpoint.
+
+    Returns:
+        ChatOpenAI: A LangChain ChatOpenAI instance configured for the API.
+    """
+    print(f"=== CONFIGURING LLM CLIENT FOR API: {api_base_url} ===")
+
+    
+    llm = ChatOpenAI(
+        openai_api_base=api_base_url,
+        openai_api_key="dummy-key", # Required by LangChain, but not used by LM Studio
+        temperature=0.7,
+        max_tokens=128,
+    )
+    return llm
+
+    # --- The following code was commented out or unreachable in the original notebook ---
+    # Example code (replace with appropriate code for your model):
+    # tokenizer = AutoTokenizer.from_pretrained(model_name)
+    # model = AutoModelForCausalLM.from_pretrained(model_name)
+    # emotion_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+    # input_emotion = "excited"
+    # input_situation = text # 'text' variable was not defined here in the original notebook
+
+    #     # Format the user message content
+    # user_content = f"Emotion: {input_emotion}\nSituation: {input_situation}"
+    #   # Create the messages list in the standard OpenAI/chat format
+    # messages = [
+    #         # Note: llama-cpp might not explicitly use a system prompt unless provided here
+    #         # or baked into the chat_format handler. You might need to add:
+    #         # {"role": "system", "content": "You are an empathetic assistant."},
+    #         {"role": "user", "content": user_content},
+    #     ]
+
+    #     # --- 3. Generate the response using create_chat_completion -- This method doesn't exist on ChatOpenAI, use invoke instead ---
+    # print("Generating response...")
+    # try:
+    #     response = llm.create_chat_completion( # This should be llm.invoke(messages)
+    #     messages=messages,
+    #     max_tokens=128,       # Max length of the generated response (adjust as needed)
+    #     temperature=0.7,    # Controls randomness (adjust)
+    #     # top_p=0.9,          # Optional: Nucleus sampling
+    #     # top_k=40,           # Optional: Top-k sampling
+    #     stop=["<|eot_id|>"], # Crucial: Stop generation when the model outputs the end-of-turn token
+    #     stream=False,        # Set to True to get token-by-token output (like TextStreamer)
+    #     )
+
+    #         # --- 4. Extract and print the response -- Access response.content with invoke ---
+    #     if response and 'choices' in response and len(response['choices']) > 0:
+    #           assistant_message = response['choices'][0]['message']['content']
+    #           print("\nAssistant Response:")
+    #           print(assistant_message.strip())
+    #           print("returning:", assistant_message.strip())
+    #           return assistant_message.strip()
+    #     else:
+    #           print("\nNo response generated or unexpected format.")
+    #           print("Full response:", response)
+
+    #           return ""
+
+    # except Exception as e:
+    #     print(f"\nAn error occurred during generation: {e}")
+    #     return ""
+    # --- End of commented out/unreachable code ---
+
+
+def load_generation_model():
+    """Load the specified Ollama model using LangChain."""
+    model_name = "meta-llama/Llama-3.2-3B-Instruct"  # ~1.1B parameters
+    
+    import os
+    token = os.getenv('HF_TOKEN'))
+    
+    try:
+        # First try loading with 4-bit quantization if bitsandbytes is available
+        from bitsandbytes.cuda_setup.main import get_compute_capability
+        
+        # Configuration for 4-bit quantization to reduce memory usage
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_quant_type="nf4",
+        )
+        
+        # Load model and tokenizer with quantization
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            quantization_config=quantization_config,
+            device_map="auto",
+            torch_dtype=torch.float16,
+            token = token
+        )
+        print("Model loaded with 4-bit quantization")
+        
+    except (ImportError, ModuleNotFoundError, PackageNotFoundError) as e:
+        print(f"Quantization not available: {e}")
+        print("Loading model in fp16 precision without quantization")
+        
+        # Fallback to fp16 without quantization
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            device_map="cuda:1",
+            torch_dtype=torch.float16,
+            low_cpu_mem_usage=True,
+            
+        )
+    
+    # Create text generation pipeline
+    generation_pipeline = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=512,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.95,
+        repetition_penalty=1.1,
+        pad_token_id=tokenizer.eos_token_id
+    )
+    
+    # Create LangChain wrapper
+    llm = HuggingFacePipeline(pipeline=generation_pipeline)
+    
+    return llm
+
+def create_prompt_templates():
+    """Create prompt templates for the assistant"""
+
+    template = """
+<|system|>
+You are an AAC (Augmentative and Alternative Communication) user (Elliot) engaging in a conversation. Your responses must reflect factual details provided in your personal context, be empathetic as guided by the emotion analysis, and align naturally with your previous chat history. You will respond directly as the AAC user, speaking in the first person (using "I", "my", "me").
+
+**Instructions:**
+1. Understand the question asked by the conversation partner.
+2. Use the provided "Context" to include accurate personal details about your life (Elliot).
+3. Reflect the empathetic tone described in the "Empathetic Response Guidance".
+4. Ensure your response fits logically within the "Chat History".
+5. Keep your response concise, empathetic, and natural.
+6. Ignore the empathetic tone described in the "Empathetic Response Guidance" if it is not related to the conversation.
+
+**Context:**
+{context}
+
+**Chat History:**
+{chat_history}
+
+**Empathetic Response Guidance:**
+{emotion_analysis}</s>
+<|user|>
+The conversation partner asked: "{question}"
+
 Please generate your response as the AAC user, following the instructions above.</s>
 <|assistant|>