Update app.py

app.py

@@ -1,33 +1,84 @@
 import streamlit as st
+import sqlite3
+import faiss
+import numpy as np
 from transformers import pipeline, M2M100ForConditionalGeneration, M2M100Tokenizer
-import PyMuPDF  # for PDF handling
+from sentence_transformers import SentenceTransformer
 from docx import Document
+import PyMuPDF
 import requests
 from bs4 import BeautifulSoup
-import faiss
-import numpy as np
-from sentence_transformers import SentenceTransformer
 from langdetect import detect
+import os
 
 # Initialize models and pipeline
 qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased")
 embedding_model = SentenceTransformer('distiluse-base-multilingual-cased-v1')
 
 # FAISS index setup (in-memory for this example)
-index = faiss.IndexFlatL2(512)
-doc_store = []
+dimension = 512  # Size of the embeddings
+index = faiss.IndexFlatL2(dimension)
 
 # Initialize translation model for on-the-fly translation
 tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
 model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
 
-# Function to translate text using the M2M100 model
+# SQLite Database Setup
+DB_PATH = "knowledge_base.db"
+
+def init_db():
+    """ Initialize the database and tables if they don't exist. """
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    c.execute('''
+        CREATE TABLE IF NOT EXISTS documents (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            content TEXT NOT NULL,
+            language TEXT,
+            embedding BLOB NOT NULL
+        )
+    ''')
+    conn.commit()
+    conn.close()
+
+def store_document(content, language, embedding):
+    """ Store document content, language, and embedding in the SQLite database. """
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    c.execute("INSERT INTO documents (content, language, embedding) VALUES (?, ?, ?)",
+              (content, language, embedding.tobytes()))
+    conn.commit()
+    conn.close()
+
+def load_documents():
+    """ Load all documents and embeddings from the SQLite database. """
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    c.execute("SELECT content, language, embedding FROM documents")
+    rows = c.fetchall()
+    conn.close()
+    
+    documents = []
+    embeddings = []
+    for content, language, embedding_blob in rows:
+        documents.append(content)
+        embeddings.append(np.frombuffer(embedding_blob, dtype=np.float32))
+    
+    return documents, np.array(embeddings)
+
 def translate_text(text, src_lang, tgt_lang):
+    """ Translate text using the M2M100 model. """
     tokenizer.src_lang = src_lang
     encoded = tokenizer(text, return_tensors="pt")
     generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id(tgt_lang))
     return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
 
+# Initialize database and FAISS index
+init_db()
+documents, embeddings = load_documents()
+if len(embeddings) > 0:
+    index.add(embeddings)
+
 # Sidebar for navigation
 st.sidebar.title("Navigation")
 page = st.sidebar.radio("Go to", ["Upload Knowledge", "Q&A"])
@@ -57,6 +108,15 @@ if page == "Upload Knowledge":
             detected_lang = detect(text)
             st.write(f"Detected language: {detected_lang}")
 
+            # Generate embeddings
+            embedding = embedding_model.encode([text])[0]
+
+            # Store the document and embedding in the database
+            store_document(text, detected_lang, embedding)
+
+            # Add the embedding to FAISS index
+            index.add(np.array([embedding], dtype=np.float32))
+            documents.append(text)
             texts.append(text)
 
         # Process URL
@@ -64,12 +124,20 @@ if page == "Upload Knowledge":
             response = requests.get(url)
             soup = BeautifulSoup(response.text, 'html.parser')
             text = soup.get_text()
+            detected_lang = detect(text)
+            st.write(f"Detected language: {detected_lang}")
+
+            # Generate embedding
+            embedding = embedding_model.encode([text])[0]
+
+            # Store the document and embedding in the database
+            store_document(text, detected_lang, embedding)
+
+            # Add the embedding to FAISS index
+            index.add(np.array([embedding], dtype=np.float32))
+            documents.append(text)
             texts.append(text)
 
-        # Create embeddings and store in FAISS
-        embeddings = embedding_model.encode(texts)
-        index.add(embeddings)
-        doc_store.extend(texts)
         st.write("Data processed and added to knowledge base!")
 
         # Provide a summary of the uploaded content
@@ -91,8 +159,8 @@ elif page == "Q&A":
             user_query = translate_text(user_query, detected_query_lang, "en")
 
         query_embedding = embedding_model.encode([user_query])
-        D, I = index.search(query_embedding, k=5)  # Retrieve top 5 documents
-        context = " ".join([doc_store[i] for i in I[0]])
+        D, I = index.search(np.array(query_embedding, dtype=np.float32), k=5)  # Retrieve top 5 documents
+        context = " ".join([documents[i] for i in I[0]])
 
         # Pass translated query and context to the QA pipeline
         result = qa_pipeline(question=user_query, context=context)