Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import spaces
 import torch
 import os
 import numpy as np
@@ -7,10 +8,11 @@ from transformers import AutoModel, AutoTokenizer
 from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
 from parler_tts import ParlerTTSForConditionalGeneration
 import soundfile as sf
-from llama_index import LLMPredictor, ServiceContext
-from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
-from llama_index.langchain_helpers.text_splitter import RecursiveCharacterTextSplitter
-from langchain import OpenAI
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.vectorstores import Chroma
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.chains import RetrievalQA
+from langchain.llms import OpenAI
 from PIL import Image
 from decord import VideoReader, cpu
 import requests
@@ -61,21 +63,21 @@ def numpy_calculate(code: str) -> str:
         return str(result)
     except Exception as e:
         return f"An error occurred: {str(e)}"
-    
+
 # Function to handle different input types
 def handle_input(user_prompt, image=None, video=None, audio=None, doc=None):
     messages = [{"role": "user", "content": user_prompt}]
-    
+
     if audio:
         transcription = client.audio.transcriptions.create(
             file=(audio.name, audio.read()),
             model="whisper-large-v3"
         )
         user_prompt = transcription.text
-    
+
     if doc:
-        index = create_rag_index(doc.name, doc.read())
-        response = index.query(user_prompt)
+        # RAG with Langchain
+        response = use_langchain_rag(doc.name, doc.read(), user_prompt)
     elif image and not video:
         image = Image.open(image).convert('RGB')
         messages[0]['content'] = [image, user_prompt]
@@ -90,15 +92,24 @@ def handle_input(user_prompt, image=None, video=None, audio=None, doc=None):
             messages=messages,
             tools=initialize_tools()
         ).choices[0].message.content
-    
+
     return response
 
-# Function to create RAG index using LlamaIndex or Langchain
-def create_rag_index(file_name, file_content):
-    docs = SimpleDirectoryReader(file_name, file_content).load_data()
-    service_context = ServiceContext.from_defaults(llm_predictor=LLMPredictor(llm=OpenAI(temperature=0)))
-    index = GPTVectorStoreIndex.from_documents(docs, service_context=service_context)
-    return index
+# Function to use Langchain for RAG
+def use_langchain_rag(file_name, file_content, query):
+    # Split the document into chunks
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
+    docs = text_splitter.create_documents([file_content])
+
+    # Create embeddings and store in the vector database
+    embeddings = OpenAIEmbeddings()
+    db = Chroma.from_documents(docs, embeddings, persist_directory=".chroma_db")  # Use a persistent directory
+
+    # Create a question-answering chain
+    qa = RetrievalQA.from_chain_type(llm=OpenAI(), chain_type="stuff", retriever=db.as_retriever())
+
+    # Get the answer
+    return qa.run(query)
 
 # Function to encode video
 def encode_video(video_path):