Update main.py

main.py

@@ -1,228 +1,160 @@
 import os  
 import shutil  
+import torch  
 from flask import Flask, render_template, request, jsonify  
-from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate, Settings  
-from llama_index.llms.huggingface import HuggingFaceInferenceAPI  
-from llama_index.embeddings.huggingface import HuggingFaceEmbedding  
-from huggingface_hub import InferenceClient  
-from transformers import AutoTokenizer, AutoModel
-from deep_translator import GoogleTranslator
-
-
-# Ensure HF_TOKEN is set  
-HF_TOKEN = os.getenv("HF_TOKEN")  
-if not HF_TOKEN:  
-    raise ValueError("HF_TOKEN environment variable not set.")  
-
-repo_id = "meta-llama/Meta-Llama-3-8B-Instruct"  
-llm_client = InferenceClient(  
-    model=repo_id,  
-    token=HF_TOKEN,  
-)  
-
-# Configure Llama index settings  
-Settings.llm = HuggingFaceInferenceAPI(  
-    model_name=repo_id,  
-    tokenizer_name=repo_id,  
-    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"  
-# )  
-# Replace the embedding model with XLM-R
-# 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"
-)
-
-# Configure tokenizer and model if required
-tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
-model = AutoModel.from_pretrained("xlm-roberta-base")
+from whoosh.index import create_in  
+from whoosh.fields import Schema, TEXT  
+from whoosh.qparser import QueryParser  
+from transformers import AutoTokenizer, AutoModel  
+from deep_translator import GoogleTranslator  
 
+# Ensure the necessary directories exist  
 PERSIST_DIR = "db"  
 PDF_DIRECTORY = 'data'  
-
-# Ensure directories exist  
 os.makedirs(PDF_DIRECTORY, exist_ok=True)  
 os.makedirs(PERSIST_DIR, exist_ok=True)  
-chat_history = []  
-current_chat_history = []  
 
+# Load the XLM-R tokenizer and model  
+tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")  
+model = AutoModel.from_pretrained("xlm-roberta-base")  
+
+# Setup Whoosh schema for indexing  
+schema = Schema(title=TEXT(stored=True), content=TEXT(stored=True))  
+
+# Create an index in the persist directory  
+if not os.path.exists(PERSIST_DIR):  
+    os.mkdir(PERSIST_DIR)  
+index = create_in(PERSIST_DIR, schema)  
+
+# Function to load documents from a directory  
+def load_documents():  
+    documents = []  
+    for filename in os.listdir(PDF_DIRECTORY):  
+        if filename.endswith(".txt"):  # Assuming documents are in .txt format  
+            with open(os.path.join(PDF_DIRECTORY, filename), 'r', encoding='utf-8') as file:  
+                content = file.read()  
+                documents.append({'title': filename, 'content': content})  
+    return documents  
+
+# Function to index documents  
+def index_documents(documents):  
+    writer = index.writer()  
+    for doc in documents:  
+        writer.add_document(title=doc['title'], content=doc['content'])  
+    writer.commit()  
+
+# Data ingestion from the directory  
 def data_ingestion_from_directory():  
     # Clear previous data by removing the persist directory  
     if os.path.exists(PERSIST_DIR):  
         shutil.rmtree(PERSIST_DIR)  # Remove the persist directory and all its contents  
-    
+
     # Recreate the persist directory after removal  
     os.makedirs(PERSIST_DIR, exist_ok=True)  
-    
-    # Load new documents from the directory  
-    new_documents = SimpleDirectoryReader(PDF_DIRECTORY).load_data()  
-    
-    # Create a new index with the new documents  
-    index = VectorStoreIndex.from_documents(new_documents)  
-    
-    # Persist the new index  
-    index.storage_context.persist(persist_dir=PERSIST_DIR)  
-
-# def handle_query(query):  
-#     context_str = ""  
-    
-#     # Build context from current chat history  
-#     for past_query, response in reversed(current_chat_history):  
-#         if past_query.strip():  
-#             context_str += f"User asked: '{past_query}'\nBot answered: '{response}'\n"  
-    
-#     chat_text_qa_msgs = [
-#         (
-#             "user",
-#             """
-#             You are the Taj Hotel voice chatbot and your name is Taj hotel helper. Your goal is to provide accurate, professional, and helpful answers to user queries based on the Taj hotel data. Always ensure your responses are clear and concise. Give response within 10-15 words only. You need to give an answer in the same language used by the user.       
-#             {context_str}
-#             Question:
-#             {query_str}
-#             """
-#         )
-#     ]
-
-
-    
-#     text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)  
-    
-#     storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)  
-#     index = load_index_from_storage(storage_context)  
-#     # context_str = ""  
-    
-#     # # Build context from current chat history  
-#     # for past_query, response in reversed(current_chat_history):  
-#     #     if past_query.strip():  
-#     #         context_str += f"User asked: '{past_query}'\nBot answered: '{response}'\n"  
-
-#     query_engine = index.as_query_engine(text_qa_template=text_qa_template, context_str=context_str)  
-#     print(f"Querying: {query}")  
-#     answer = query_engine.query(query)  
-
-#     # Extracting the response  
-#     if hasattr(answer, 'response'):  
-#         response = answer.response  
-#     elif isinstance(answer, dict) and 'response' in answer:  
-#         response = answer['response']  
-#     else:  
-#         response = "I'm sorry, I couldn't find an answer to that."  
-
-#     # Append to chat history  
-#     current_chat_history.append((query, response))  
-#     return response
-def handle_query(query):
-    chat_text_qa_msgs = [
-        (
-            "user",
-            """
-            You are the Hotel voice chatbot and your name is hotel helper. Your goal is to provide accurate, professional, and helpful answers to user queries based on the hotel's data. Always ensure your responses are clear and concise. Give response within 10-15 words only. You need to give an answer in the same language used by the user.       
-            {context_str}
-            Question:
-            {query_str}
-            """
-        )
-    ]
-    text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
-    
-    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
-    index = load_index_from_storage(storage_context)
-    context_str = ""
-    for past_query, response in reversed(current_chat_history):
-        if past_query.strip():
-            context_str += f"User asked: '{past_query}'\nBot answered: '{response}'\n"
-
-    query_engine = index.as_query_engine(text_qa_template=text_qa_template, context_str=context_str)
-    print(query)
-    answer = query_engine.query(query)
-
-    if hasattr(answer, 'response'):
-        response = answer.response
-    elif isinstance(answer, dict) and 'response' in answer:
-        response = answer['response']
-    else:
-        response = "Sorry, I couldn't find an answer."
-    current_chat_history.append((query, response))
-    return response
 
+    # Load new documents from the directory  
+    new_documents = load_documents()  
+
+    # Index the new documents  
+    index_documents(new_documents)  
+
+# Function to retrieve documents based on a query  
+def retrieve_documents(query):  
+    with index.searcher() as searcher:  
+        query_parser = QueryParser("content", index.schema)  
+        query_object = query_parser.parse(query)  
+        results = searcher.search(query_object)  
+        return [(result['title'], result['content']) for result in results]  
+
+# Function to generate embeddings  
+def get_embeddings(text):  
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)  
+    with torch.no_grad():  
+        outputs = model(**inputs)  
+    embeddings = outputs.last_hidden_state.mean(dim=1)  # Average pooling  
+    return embeddings.squeeze().numpy()  
+
+# Function to handle queries and generate responses  
+def handle_query(query):  
+    retrieved_docs = retrieve_documents(query)  
+
+    if not retrieved_docs:  
+        return "Sorry, I couldn't find an answer."  
+
+    # Construct a response using the retrieved documents  
+    response = "Here are some insights based on your query:\n" + "\n".join(  
+        [f"Title: {title}\nContent: {content[:100]}..." for title, content in retrieved_docs]  
+    )  
+    return response  
+
+# Initialize Flask app  
 app = Flask(__name__)  
 
 # Data ingestion  
 data_ingestion_from_directory()  
 
-# Generate Response
+# Generate Response  
 def generate_response(query, language):  
     try:  
         # Call the handle_query function to get the response  
-        bot_response = handle_query(query)
-        
-        # Map of supported languages
-        supported_languages = {
-            "hindi": "hi",
-            "bengali": "bn",
-            "telugu": "te",
-            "marathi": "mr",
-            "tamil": "ta",
-            "gujarati": "gu",
-            "kannada": "kn",
-            "malayalam": "ml",
-            "punjabi": "pa",
-            "odia": "or",
-            "urdu": "ur",
-            "assamese": "as",
-            "sanskrit": "sa",
-            "arabic": "ar",
-            "australian": "en-AU",
-            "bangla-india": "bn-IN",
-            "chinese": "zh-CN",
-            "dutch": "nl",
-            "french": "fr",
-            "filipino": "tl",
-            "greek": "el",
-            "indonesian": "id",
-            "italian": "it",
-            "japanese": "ja",
-            "korean": "ko",
-            "latin": "la",
-            "nepali": "ne",
-            "portuguese": "pt",
-            "romanian": "ro",
-            "russian": "ru",
-            "spanish": "es",
-            "swedish": "sv",
-            "thai": "th",
-            "ukrainian": "uk",
-            "turkish": "tr"
-        }
-        
-        # Initialize the translated text
-        translated_text = bot_response
-        
-        # Translate only if the language is supported and not English
-        try:
-            if language in supported_languages:
-                target_lang = supported_languages[language]
-                translated_text = GoogleTranslator(source='en', target=target_lang).translate(bot_response)
-                print(translated_text)
-            else:
-                print(f"Unsupported language: {language}")
-        except Exception as e:
-            # Handle translation errors
-            print(f"Translation error: {e}")
-            translated_text = "Sorry, I couldn't translate the response."
-        
-        # Append to chat history
-        chat_history.append((query, translated_text))
+        bot_response = handle_query(query)  
+
+        # Map of supported languages  
+        supported_languages = {  
+            "hindi": "hi",  
+            "bengali": "bn",  
+            "telugu": "te",  
+            "marathi": "mr",  
+            "tamil": "ta",  
+            "gujarati": "gu",  
+            "kannada": "kn",  
+            "malayalam": "ml",  
+            "punjabi": "pa",  
+            "odia": "or",  
+            "urdu": "ur",  
+            "assamese": "as",  
+            "sanskrit": "sa",  
+            "arabic": "ar",  
+            "australian": "en-AU",  
+            "bangla-india": "bn-IN",  
+            "chinese": "zh-CN",  
+            "dutch": "nl",  
+            "french": "fr",  
+            "filipino": "tl",  
+            "greek": "el",  
+            "indonesian": "id",  
+            "italian": "it",  
+            "japanese": "ja",  
+            "korean": "ko",  
+            "latin": "la",  
+            "nepali": "ne",  
+            "portuguese": "pt",  
+            "romanian": "ro",  
+            "russian": "ru",  
+            "spanish": "es",  
+            "swedish": "sv",  
+            "thai": "th",  
+            "ukrainian": "uk",  
+            "turkish": "tr"  
+        }  
+
+        # Initialize the translated text  
+        translated_text = bot_response  
+
+        # Translate only if the language is supported and not English  
+        try:  
+            if language in supported_languages:  
+                target_lang = supported_languages[language]  
+                translated_text = GoogleTranslator(source='auto', target=target_lang).translate(bot_response)  
+            else:  
+                print(f"Unsupported language: {language}")  
+        except Exception as e:  
+            print(f"Translation error: {e}")  
+            translated_text = "Sorry, I couldn't translate the response."  
+
         return translated_text  
     except Exception as e:  
-        return f"Error fetching the response: {str(e)}"
+        return f"Error fetching the response: {str(e)}"  
 
 # Route for the homepage  
 @app.route('/')  
@@ -233,12 +165,12 @@ def index():
 @app.route('/chat', methods=['POST'])  
 def chat():  
     try:  
-        user_message = request.json.get("message")
-        language = request.json.get("language")
+        user_message = request.json.get("message")  
+        language = request.json.get("language")  
         if not user_message:  
             return jsonify({"response": "Please say something!"})  
 
-        bot_response = generate_response(user_message,language)  
+        bot_response = generate_response(user_message, language)  
         return jsonify({"response": bot_response})  
     except Exception as e:  
         return jsonify({"response": f"An error occurred: {str(e)}"})