bunch of changes

app.py

@@ -1,4 +1,200 @@
 import streamlit as st
+from docx import Document
+import os
+from langchain_core.prompts import PromptTemplate
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import time
+from sentence_transformers import SentenceTransformer
+from langchain.vectorstores import Chroma
+from langchain.docstore.document import Document
+from langchain_community.embeddings import HuggingFaceEmbeddings
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+
+docs_folder = "./converted_docs"
+
+
+# Function to load .docx files from Google Drive folder
+def load_docx_files_from_drive(drive_folder):
+    docx_files = [f for f in os.listdir(drive_folder) if f.endswith(".docx")]
+    documents = []
+
+    for file_name in docx_files:
+        file_path = os.path.join(drive_folder, file_name)
+        doc = Document(file_path)
+        content = "\n".join([p.text for p in doc.paragraphs if p.text.strip()])
+        documents.append(content)
+
+    return documents
+
+
+# Load .docx files from Google Drive folder
+documents = load_docx_files_from_drive(docs_folder)
+
+
+def split_extracted_text_into_chunks(documents):
+    # List to hold all chunks
+    chunks = []
+
+    for doc_text in documents:
+        # Split the document text into lines
+        lines = doc_text.splitlines()
+
+        # Initialize variables for splitting
+        current_chunk = []
+        for line in lines:
+            # Check if the line starts with "File Name:"
+            if line.startswith("File Name:"):
+                # If there's a current chunk, save it before starting a new one
+                if current_chunk:
+                    chunks.append("\n".join(current_chunk))
+                    current_chunk = []  # Reset the current chunk
+
+            # Add the line to the current chunk
+            current_chunk.append(line)
+
+        # Add the last chunk for the current document
+        if current_chunk:
+            chunks.append("\n".join(current_chunk))
+
+    return chunks
+
+
+# Split the extracted documents into chunks
+chunks = split_extracted_text_into_chunks(documents)
+
+
+def save_chunks_to_file(chunks, output_file_path):
+    # Open the file in write mode
+    with open(output_file_path, "w", encoding="utf-8") as file:
+        for i, chunk in enumerate(chunks, start=1):
+            # Write each chunk with a header for easy identification
+            file.write(f"Chunk {i}:\n")
+            file.write(chunk)
+            file.write("\n" + "=" * 50 + "\n")
+
+
+# Path to save the chunks file
+output_file_path = "./chunks_output.txt"
+
+# Split the extracted documents into chunks
+chunks = split_extracted_text_into_chunks(documents)
+
+# Save the chunks to the file
+save_chunks_to_file(chunks, output_file_path)
+
+
+# Step 1: Load the model through LangChain's wrapper
+embedding_model = HuggingFaceEmbeddings(
+    model_name="Omartificial-Intelligence-Space/Arabic-Triplet-Matryoshka-V2"
+)
+
+
+# Step 2: Embed the chunks (now simplified)
+def embed_chunks(chunks):
+    return [
+        {"chunk": chunk, "embedding": embedding_model.embed_query(chunk)}
+        for chunk in chunks
+    ]
+
+
+embeddings = embed_chunks(chunks)
+
+
+# Step 3: Prepare documents (unchanged)
+def prepare_documents_for_chroma(embeddings):
+    return [
+        Document(page_content=entry["chunk"], metadata={"chunk_index": i})
+        for i, entry in enumerate(embeddings, start=1)
+    ]
+
+
+documents = prepare_documents_for_chroma(embeddings)
+
+# Step 4: Create Chroma store (fixed)
+vectorstore = Chroma.from_documents(
+    documents=documents,
+    embedding=embedding_model,  # Proper embedding object
+    persist_directory="./chroma_db",  # Optional persistence
+)
+
+
+class RAGPipeline:
+    def __init__(self, vectorstore, model_name="CohereForAI/aya-expanse-8b", k=6):
+        self.vectorstore = vectorstore
+        self.model_name = model_name
+        self.k = k
+        self.retriever = self.vectorstore.as_retriever(
+            search_type="mmr", search_kwargs={"k": self.k}
+        )
+        self.prompt_template = PromptTemplate.from_template(self._get_template())
+
+        # Load model and tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_name, torch_dtype=torch.bfloat16, device_map="auto"
+        )
+
+    def _get_template(self):
+        return """\
+        <s>[INST] <<SYS>>
+        أنت مساعد مفيد يقدم إجابات باللغة العربية بناءً على السياق المقدم.
+        - أجب فقط باللغة العربية
+        - إذا لم تجد إجابة في السياق، قل أنك لا تعرف
+        - كن دقيقاً وواضحاً في إجاباتك
+        <</SYS>>
+
+        السياق: {context}
+
+        السؤال: {question}
+        الإجابة: [/INST]\
+        """
+
+    def generate_response(self, question):
+        retrieved_docs = self._retrieve_documents(question)
+        prompt = self._create_prompt(retrieved_docs, question)
+        response = self._generate_response(prompt)
+        return response
+
+    def _retrieve_documents(self, question):
+        start = time.time()
+        retrieved_docs = self.retriever.invoke(question)
+        result = {f"doc_{i}": doc.page_content for i, doc in enumerate(retrieved_docs)}
+        end = time.time()
+        time_lapsed = end - start
+        print(f"Time lapsed in Retreival: {time_lapsed}")
+        return result
+
+    def _create_prompt(self, docs, question):
+        return self.prompt_template.format(context=docs, question=question)
+
+    def _generate_response(self, prompt):
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+
+        start = time.time()
+        outputs = self.model.generate(
+            inputs.input_ids,
+            max_new_tokens=1024,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=self.tokenizer.eos_token_id,
+        )
+        end = time.time()
+        time_lapsed = end - start
+        print(f"Time lapsed in Generation: {time_lapsed}")
+
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Extract only the assistant's response after [/INST]
+        return response.split("[/INST]")[-1].strip()
+
+
+rag_pipeline = RAGPipeline(vectorstore)
+
+
+question = st.text_area("أدخل سؤالك هنا")
+if st.button("Generate Answer"):
+    response = rag_pipeline.generate_response(question)
+    st.write(response)
+    print("Question: ", question)
+    print("Response: ", response)

requirements.txt

@@ -0,0 +1,6 @@
+torch
+python-docx
+langchain
+langchain_community==0.0.36
+langchain_chroma
+sentence-transformers