Update app.py

app.py

@@ -4,42 +4,25 @@ import os
 from gtts import gTTS
 from deep_translator import GoogleTranslator
 import logging
-from llama_index import VectorStoreIndex, Document, SimpleDirectoryReader
-from llama_index.node_parser import SimpleNodeParser
-from llama_index.llms import HuggingFaceLLM
-from llama_index import ServiceContext, set_global_service_context
-import torch
+from llama_index import VectorStoreIndex, Document
+from llama_index.embeddings import HuggingFaceEmbedding
+from llama_index import ServiceContext
+from groq import Groq
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s | %(levelname)s | %(message)s')
 
-# Initialize the LLM with a smaller context window
-try:
-    llm = HuggingFaceLLM(
-        context_window=512,  # Reduced from 1024
-        max_new_tokens=256,
-        generate_kwargs={"temperature": 0.7, "do_sample": False},
-        tokenizer_name="gpt2",
-        model_name="gpt2",
-        device_map="auto",
-        tokenizer_kwargs={"max_length": 512},  # Reduced from 1024
-        model_kwargs={"torch_dtype": torch.float32},
-    )
-except ImportError:
-    # Fallback if Accelerate is not available
-    llm = HuggingFaceLLM(
-        context_window=512,  # Reduced from 1024
-        max_new_tokens=256,
-        generate_kwargs={"temperature": 0.7, "do_sample": False},
-        tokenizer_name="gpt2",
-        model_name="gpt2",
-        tokenizer_kwargs={"max_length": 512},  # Reduced from 1024
-        model_kwargs={"torch_dtype": torch.float32},
-    )
+# Initialize Groq client
+groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))
+
+# Initialize the embedding model
+embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
-# Initialize the ServiceContext with a chunk size
-node_parser = SimpleNodeParser.from_defaults(chunk_size=256)  # Adjust chunk size as needed
-service_context = ServiceContext.from_defaults(llm=llm, embed_model="local", node_parser=node_parser)
-set_global_service_context(service_context)
+# Initialize the ServiceContext
+service_context = ServiceContext.from_defaults(embed_model=embed_model)
 
 # Initialize the index
 index = None
@@ -73,28 +56,40 @@ audio_language_dict = {
 def index_text(text: str) -> str:
     global index
     try:
-        documents = [Document(text=text)]
+        document = Document(text=text)
         if index is None:
-            index = VectorStoreIndex.from_documents(documents)
+            index = VectorStoreIndex.from_documents([document], service_context=service_context)
         else:
-            index.insert(documents[0])
+            index.insert(document)
         return "Text indexed successfully."
     except Exception as e:
         logging.error(f"Error in indexing: {str(e)}")
         return f"Error indexing text: {str(e)}"
 
-def chat_with_context(question: str) -> str:
+def chat_with_context(question: str, model: str) -> str:
     global index
     if index is None:
         return "Please index some text first."
     
     try:
         query_engine = index.as_query_engine(
-            similarity_top_k=2,  # Adjust as needed
+            similarity_top_k=2,
             response_mode="compact"
         )
-        response = query_engine.query(question)
-        return str(response)
+        context = query_engine.query(question).response
+        
+        prompt = f"Context: {context}\n\nQuestion: {question}\n\nAnswer:"
+        
+        chat_completion = groq_client.chat.completions.create(
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt,
+                }
+            ],
+            model=model,
+        )
+        return chat_completion.choices[0].message.content
     except Exception as e:
         logging.error(f"Error in chat: {str(e)}")
         return f"Error in chat: {str(e)}"
@@ -140,6 +135,11 @@ with gr.Blocks() as iface:
     chat_group = gr.Group(visible=False)
     with chat_group:
         chat_input = gr.Textbox(label="Ask a question about the indexed text")
+        chat_model = gr.Dropdown(
+            choices=["llama3-70b-8192", "mixtral-8x7b-32768", "gemma-7b-it"],
+            label="Select Chat Model",
+            value="llama3-70b-8192"
+        )
         chat_button = gr.Button("Ask")
         chat_output = gr.Textbox(label="Answer", interactive=False)
     
@@ -171,7 +171,7 @@ with gr.Blocks() as iface:
     convert_button.click(convert_text, inputs=[text_input, translation_lang_dropdown], outputs=translated_text)
     index_button.click(index_text, inputs=[translated_text], outputs=[index_status])
     use_chat.change(update_chat_visibility, inputs=[use_chat], outputs=[chat_group])
-    chat_button.click(chat_with_context, inputs=[chat_input], outputs=[chat_output])
+    chat_button.click(chat_with_context, inputs=[chat_input, chat_model], outputs=[chat_output])
     
     generate_button.click(
         generate_speech,