RAG LlamaIndex

app.py

@@ -1,126 +1,156 @@
-import os
 import gradio as gr
-from argparse import ArgumentParser
+import os
+import warnings
+import asyncio
+from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Document, Settings
+from llama_index.llms.cerebras import Cerebras
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from groq import Groq
-import base64
 import io
 
-# Initialize Groq client
-API_KEY = os.environ['GROQ_API_KEY']
-client = Groq(api_key=API_KEY)
-
-REVISION = 'v1.0.4'
-
-def _get_args():
-    parser = ArgumentParser()
-    parser.add_argument("--revision", type=str, default=REVISION)
-    parser.add_argument("--share", action="store_true", default=False, help="Create a publicly shareable link for the interface.")
-    return parser.parse_args()
-
-def process_image(image):
-    buffered = io.BytesIO()
-    image.save(buffered, format="JPEG")
-    return buffered.getvalue()
-
-def translate_audio(audio_file):
-    with open(audio_file, "rb") as file:
-        translation = client.audio.translations.create(
-            file=(audio_file, file.read()),
-            model="whisper-large-v3",
-            response_format="json",
-            temperature=0.0
-        )
-    return translation.text
-
-def transcribe_audio(audio_file):
-    with open(audio_file, "rb") as file:
-        transcription = client.audio.transcriptions.create(
-            file=(audio_file, file.read()),
-            model="whisper-large-v3",
-            response_format="json",
-            temperature=0.0
-        )
-    return transcription.text
-
-def predict(chat_history, query, image, audio, translate):
-    final_query = query.strip()
-    
-    if audio:
-        audio_file_path = audio
-        if translate:
-            translation_text = translate_audio(audio_file_path)
-            final_query = translation_text.strip()
-            chat_history.append({'role': 'assistant', 'content': translation_text})
-        else:
-            transcribed_text = transcribe_audio(audio_file_path)
-            final_query = f"{final_query} {transcribed_text}".strip()
+# Suppress warnings
+warnings.filterwarnings("ignore", message=".*clean_up_tokenization_spaces.*")
+
+# Global variables
+index = None
+query_engine = None
+
+# Load Cerebras API key from Hugging Face secrets
+api_key = os.getenv("CEREBRAS_API_KEY")
+if not api_key:
+    raise ValueError("CEREBRAS_API_KEY is not set in Hugging Face Secrets.")
+else:
+    print("Cerebras API key loaded successfully.")
+
+# Initialize Cerebras LLM and embedding model
+os.environ["CEREBRAS_API_KEY"] = api_key
+llm = Cerebras(model="llama3.1-70b", api_key=os.environ["CEREBRAS_API_KEY"])
+Settings.llm = llm
+embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+# Initialize Groq client for Whisper Large V3
+groq_api_key = os.getenv("GROQ_API_KEY")
+if not groq_api_key:
+    raise ValueError("GROQ_API_KEY is not set.")
+else:
+    print("Groq API key loaded successfully.")
+client = Groq(api_key=groq_api_key)
+
+# Function for audio transcription and translation (Whisper Large V3 from Groq)
+def transcribe_or_translate_audio(audio_file, translate=False):
+    """
+    Transcribes or translates audio using Whisper Large V3 via Groq API.
+    """
+    try:
+        with open(audio_file, "rb") as file:
+            if translate:
+                result = client.audio.translations.create(
+                    file=(audio_file, file.read()),
+                    model="whisper-large-v3",
+                    response_format="json",
+                    temperature=0.0
+                )
+                return result.text
+            else:
+                result = client.audio.transcriptions.create(
+                    file=(audio_file, file.read()),
+                    model="whisper-large-v3",
+                    response_format="json",
+                    temperature=0.0
+                )
+                return result.text
+    except Exception as e:
+        return f"Error processing audio: {str(e)}"
 
-    image_data = process_image(image) if image else None
-    messages = create_messages(final_query, image_data)
+# Function to load documents and create index
+def load_documents(file_objs):
+    global index, query_engine
+    try:
+        if not file_objs:
+            return "Error: No files selected."
+
+        documents = []
+        document_names = []
+        for file_obj in file_objs:
+            document_names.append(file_obj.name)
+            loaded_docs = SimpleDirectoryReader(input_files=[file_obj.name]).load_data()
+            for doc in loaded_docs:
+                doc.metadata["source"] = file_obj.name
+                documents.append(doc)
+
+        if not documents:
+            return "No documents found in the selected files."
 
-    if not messages:
-        error_message = "No valid input provided. Please enter a query or upload an image/audio."
-        chat_history.append({'role': 'assistant', 'content': error_message})
-        return chat_history
+        index = VectorStoreIndex.from_documents(documents, llm=llm, embed_model=embed_model)
+        query_engine = index.as_query_engine()
+
+        return f"Successfully loaded {len(documents)} documents from the files: {', '.join(document_names)}"
+    except Exception as e:
+        return f"Error loading documents: {str(e)}"
+
+async def perform_rag(query, history, audio_file=None, translate_audio=False):
+    global query_engine
+    if query_engine is None:
+        return history + [("Please load documents first.", None)]
 
     try:
-        completion = client.chat.completions.create(
-            model="llama-3.2-90b-vision-preview",
-            messages=messages,
-            temperature=1,
-            max_tokens=1500,
-            top_p=1,
-            stream=False,
-        )
-
-        response_text = completion.choices[0].message.content.strip()
-        chat_history.append({'role': 'user', 'content': final_query})
-        chat_history.append({'role': 'assistant', 'content': response_text})
+        # Handle audio input if provided
+        if audio_file:
+            transcription = transcribe_or_translate_audio(audio_file, translate=translate_audio)
+            query = f"{query} {transcription}".strip()
+
+        response = await asyncio.to_thread(query_engine.query, query)
+        answer = str(response)  # Directly get the answer from the response
+
+        # If relevant documents are available, add sources without the "Sources" label
+        if hasattr(response, "get_documents"):
+            relevant_docs = response.get_documents()
+            if relevant_docs:
+                sources = "\n\n".join([f"{doc.metadata.get('source', 'No source available')}" for doc in relevant_docs])
+            else:
+                sources = ""
+        else:
+            sources = ""
+
+        # Combine answer with sources (if any) without additional labels
+        final_result = f"{answer}\n\n{sources}".strip()
+
+        # Return updated history with the final result
+        return history + [(query, final_result)]
     except Exception as e:
-        response_text = f"Error: {str(e)}"
-        chat_history.append({'role': 'user', 'content': final_query})
-        chat_history.append({'role': 'assistant', 'content': response_text})
-    
-    return chat_history
-
-def create_messages(query, image_data):
-    messages = []
-    if query:
-        messages.append({'role': 'user', 'content': query})
-    if image_data:
-        image_base64 = f"data:image/jpeg;base64,{base64.b64encode(image_data).decode()}"
-        messages.append({
-            'role': 'user',
-            'content': [
-                {"type": "text", "text": "Please analyze this image."},
-                {"type": "image_url", "image_url": {"url": image_base64}}
-            ]
-        })
-    return messages
-
-def clear_history():
-    return []
-
-def main():
-    args = _get_args()
+        return history + [(query, f"Error processing query: {str(e)}")]
+
+# Function to clear the session and reset variables
+def clear_all():
+    global index, query_engine
+    index = None
+    query_engine = None
+    return None, "", [], ""  # Reset file input, load output, chatbot, and message input to default states
+
+# Create the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# RAG Multi-file Chat Application with Speech-to-Text")
     
-    with gr.Blocks(css="#chatbox {height: 400px; background-color: #f9f9f9; padding: 20px; border-radius: 10px; }") as demo:
-        gr.Markdown("<h1 style='text-align: center; color: #4a4a4a;'>Llama-3.2-90b-vision-preview</h1>")
-        
-        chatbox = gr.Chatbot(type='messages', elem_id="chatbox")
-        query = gr.Textbox(label="Type your query here...", placeholder="Enter your question or command...", lines=2)
-        image_input = gr.Image(type="pil", label="Upload Image")
-        audio_input = gr.Audio(type="filepath", label="Upload Audio")
-        translate_checkbox = gr.Checkbox(label="Translate Audio to English Text")
-
-        with gr.Row():
-            submit_btn = gr.Button("Submit", variant="primary", elem_id="submit-btn")
-            clear_btn = gr.Button("Clear History", variant="secondary", elem_id="clear-btn")
-        
-        submit_btn.click(predict, inputs=[chatbox, query, image_input, audio_input, translate_checkbox], outputs=chatbox)
-        clear_btn.click(clear_history, outputs=chatbox)
-        
-    demo.launch(share=args.share)
-
-if __name__ == '__main__':
-    main()
+    with gr.Row():
+        file_input = gr.File(label="Select files to load", file_count="multiple")
+        load_btn = gr.Button("Load Documents")
+        load_output = gr.Textbox(label="Load Status")
+
+    msg = gr.Textbox(label="Enter your question")
+    audio_input = gr.Audio(type="filepath", label="Upload Audio")
+    translate_checkbox = gr.Checkbox(label="Translate Audio to English Text", value=False)
+    chatbot = gr.Chatbot()
+    clear = gr.Button("Clear")
+
+    # Set up event handlers
+    load_btn.click(load_documents, inputs=[file_input], outputs=[load_output])
+
+    # Event handler for audio input to directly trigger processing and chat response
+    audio_input.change(perform_rag, inputs=[msg, chatbot, audio_input, translate_checkbox], outputs=[chatbot])
+
+    clear.click(clear_all, outputs=[file_input, load_output, chatbot, msg], queue=False)
+
+# Run the app
+if __name__ == "__main__":
+    demo.queue()
+    demo.launch()