Update main.py

main.py

@@ -1,77 +1,183 @@
 import os
-from flask import Flask, request, jsonify
-from llama_index import SimpleDirectoryReader, StorageContext, VectorStoreIndex, load_index_from_storage, ChatPromptTemplate
+import shutil
+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 transformers import AutoTokenizer, AutoModel, XLMRobertaXLForMultipleChoice
 from deep_translator import GoogleTranslator
+import torch
 
 # Ensure HF_TOKEN is set
 HF_TOKEN = os.getenv("HF_TOKEN")
 if not HF_TOKEN:
     raise ValueError("HF_TOKEN environment variable not set.")
 
-# Hugging Face model configuration
-REPO_ID = "facebook/xlm-roberta-xl"  # Use xlm-roberta-xl model
-tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
-model = AutoModel.from_pretrained(REPO_ID)
+repo_id = "facebook/xlm-roberta-xl"
+tokenizer = AutoTokenizer.from_pretrained(repo_id)
+model = XLMRobertaXLForMultipleChoice.from_pretrained(repo_id)
 
-# Flask app
-app = Flask(__name__)
-
-# Directories for storing data
 PERSIST_DIR = "db"
-PDF_DIRECTORY = "data"
+PDF_DIRECTORY = 'data'
+
+# Ensure directories exist
 os.makedirs(PDF_DIRECTORY, exist_ok=True)
 os.makedirs(PERSIST_DIR, exist_ok=True)
-
-# Initialize variables
 chat_history = []
+current_chat_history = []
 
-# Function to ingest documents
 def data_ingestion_from_directory():
+    # Clear previous data by removing the persist directory
     if os.path.exists(PERSIST_DIR):
-        os.system(f"rm -rf {PERSIST_DIR}")  # Clear previous data
+        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)
 
-    documents = SimpleDirectoryReader(PDF_DIRECTORY).load_data()
-    index = VectorStoreIndex.from_documents(documents)
+    # 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)
 
-# Function to handle queries
 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)
-    query_engine = index.as_query_engine()
+    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
 
-    chat_prompt = ChatPromptTemplate.from_messages([
-        ("user", "User asked: {query_str}"),
-        ("assistant", "Answer: {response}"),
-    ])
+def evaluate_model(prompt, choice0, choice1):
+    labels = torch.tensor(0).unsqueeze(0)  # choice0 is correct, batch size 1
+    encoding = tokenizer([prompt, prompt], [choice0, choice1], return_tensors="pt", padding=True)
+    outputs = model(**{k: v.unsqueeze(0) for k, v in encoding.items()}, labels=labels)  # batch size is 1
 
-    result = query_engine.query(query, prompt_template=chat_prompt)
-    return result.response if hasattr(result, 'response') else "No relevant answer found."
+    # the linear classifier still needs to be trained
+    loss = outputs.loss
+    logits = outputs.logits
+    return loss, logits
+
+app = Flask(__name__)
+
+# Data ingestion
+data_ingestion_from_directory()
+
+# 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))
+        return translated_text
+    except Exception as e:
+        return f"Error fetching the response: {str(e)}"
 
-# Route for homepage
-@app.route("/")
+# Route for the homepage
+@app.route('/')
 def index():
-    return "Welcome to the RAG Application using xlm-roberta-xl!"
+    return render_template('index.html')
 
 # Route to handle chatbot messages
-@app.route("/chat", methods=["POST"])
+@app.route('/chat', methods=['POST'])
 def chat():
     try:
         user_message = request.json.get("message")
+        language = request.json.get("language")
         if not user_message:
-            return jsonify({"response": "Please provide a message!"})
+            return jsonify({"response": "Please say something!"})
 
-        # Generate response
-        response = handle_query(user_message)
-        chat_history.append({"user": user_message, "bot": response})
-        return jsonify({"response": response})
+        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)}"})
 
-if __name__ == "__main__":
-    # Ingest data before starting the app
-    data_ingestion_from_directory()
+if __name__ == '__main__':
     app.run(debug=True)