Update rag_gradio_app.py

rag_gradio_app.py

@@ -1,56 +1,45 @@
 import gradio as gr
-import torch
+import openai
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
-from sentence_transformers import SentenceTransformer
-import openai
 
-# Load pre-trained embedding model
-model_name = 'intfloat/e5-small'
+# Set API Key
+openai.api_key = "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"
+
+# Load embedding model
+model_name = "intfloat/e5-small"
 embedding_model = HuggingFaceEmbeddings(model_name=model_name)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load ChromaDB
-persist_directory = './docs/chroma/'
+persist_directory = "./docs/chroma/"
 vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_model)
 
-# OpenAI API Key
-openai.api_key = 'your-api-key'
-
-def retrieve_documents(question, k=5):
-    """Retrieve top K relevant documents from ChromaDB"""
-    docs = vectordb.similarity_search(question, k=k)
-    return [doc.page_content for doc in docs]
-
-def generate_response(question, context):
-    """Generate response using OpenAI GPT-4"""
-    full_prompt = f"Context: {context}\n\nQuestion: {question}"
+# Define RAG function
+def rag_pipeline(question):
+    """Retrieve relevant documents and generate AI response"""
+    retrieved_docs = vectordb.similarity_search(question, k=5)
+    context = " ".join([doc.page_content for doc in retrieved_docs])
+    
+    # Generate AI response
+    full_prompt = f"Context: {context}\\n\\nQuestion: {question}"
     response = openai.ChatCompletion.create(
         model="gpt-4",
         messages=[{"role": "user", "content": full_prompt}],
         max_tokens=300,
         temperature=0.7
     )
-    return response['choices'][0]['message']['content'].strip()
-
-def rag_pipeline(question):
-    """Full RAG Pipeline - Retrieve Docs & Generate Response"""
-    retrieved_docs = retrieve_documents(question, k=5)
-    context = " ".join(retrieved_docs)
-    response = generate_response(question, context)
-    return response, retrieved_docs
-
-def gradio_interface(question):
-    response, retrieved_docs = rag_pipeline(question)
-    return response, "\n\n".join(retrieved_docs)
+    
+    return response['choices'][0]['message']['content'].strip(), retrieved_docs
 
-# Create Gradio App
+# Gradio UI
 iface = gr.Interface(
-    fn=gradio_interface,
+    fn=rag_pipeline,
     inputs=gr.Textbox(label="Enter your question"),
     outputs=[gr.Textbox(label="Generated Response"), gr.Textbox(label="Retrieved Documents")],
     title="RAG-Based Question Answering System",
-    description="Enter a question and retrieve relevant documents along with the AI-generated response."
+    description="Enter a question and retrieve relevant documents with AI-generated response."
 )
 
-iface.launch()
+# Launch Gradio app
+if __name__ == "__main__":
+    iface.launch()