Update app.py

app.py

@@ -1,109 +1,57 @@
 import gradio as gr
-import openai
-import os
-from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
-from langchain.schema import Document
-from sentence_transformers import SentenceTransformer
-from datasets import load_dataset
-import nltk
-import nltk
-nltk.download('punkt')  # ✅ Correct package
-import nltk
-import os
-nltk.data.path.append("/usr/local/share/nltk_data")  # ✅ Set correct path
-nltk.download("punkt")
-
+import openai
+import torch
 
-# ✅ Load the Sentence Transformer Embedding Model
-model_name = "sentence-transformers/all-MiniLM-L6-v2"
+# Initialize models and configurations
+model_name = 'intfloat/e5-small'
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 embedding_model = HuggingFaceEmbeddings(model_name=model_name)
+embedding_model.client.to(device)
 
-# ✅ Set OpenAI API Key
-openai.api_key = os.getenv("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")
-
-# ✅ Download NLTK Tokenizer
-nltk.download('punkt')
-
-# ✅ Load and Chunk Dataset
-def chunk_documents(documents, max_chunk_size=500):
-    chunks = []
-    for doc in documents:
-        sentences = nltk.sent_tokenize(doc)
-        current_chunk = ""
-        for sentence in sentences:
-            if len(current_chunk) + len(sentence) <= max_chunk_size:
-                current_chunk += sentence + " "
-            else:
-                chunks.append(current_chunk.strip())
-                current_chunk = sentence + " "
-        if current_chunk:
-            chunks.append(current_chunk.strip())
-    return chunks
-
-# ✅ Load Dataset and Prepare ChromaDB
-dataset = load_dataset("rungalileo/ragbench", "techqa")  # Example dataset
-original_documents = dataset['train']['documents']
-chunked_documents = chunk_documents(original_documents)
-
-persist_directory = "chroma_db_directory"
-documents = [Document(page_content=chunk) for chunk in chunked_documents]
-
-# ✅ Initialize ChromaDB
-vectordb = Chroma.from_documents(
-    documents=documents,
-    embedding=embedding_model,
-    persist_directory=persist_directory
+# Initialize Chroma
+vectordb = Chroma(
+    persist_directory='./docs/chroma/',
+    embedding_function=embedding_model
 )
-vectordb.persist()
-
-# ✅ Function to Retrieve Relevant Documents
-def retrieve_documents(question, k=5):
-    docs = vectordb.similarity_search(question, k=k)
-    if not docs:
-        return ["⚠️ No relevant documents found. Try a different query."]
-    return [doc.page_content for doc in docs]
-
-# ✅ Function to Generate AI Response
-def generate_response(question, context):
-    if not context or "No relevant documents found." in context:
-        return "No relevant context available. Try a different query."
-
-    full_prompt = f"Context: {context}\n\nQuestion: {question}"
 
-    try:
-        response = openai.ChatCompletion.create(
-            model="gpt-4",
-            messages=[
-                {"role": "system", "content": "You are an AI assistant that answers user queries based on the given context."},
-                {"role": "user", "content": full_prompt}
-            ],
-            max_tokens=300,
-            temperature=0.7
-        )
-        return response['choices'][0]['message']['content'].strip()
-    except Exception as e:
-        return f"Error generating response: {str(e)}"
-
-# ✅ Full RAG Pipeline
-def rag_pipeline(question):
-    retrieved_docs = retrieve_documents(question, k=5)
-    context = " ".join(retrieved_docs)
-    response = generate_response(question, context)
-    return response, "\n\n".join(retrieved_docs)
-
-# ✅ Gradio UI Interface
-iface = gr.Interface(
-    fn=rag_pipeline,
-    inputs=gr.Textbox(label="Enter your question"),
+def process_query(query):
+    # Get relevant documents
+    relevant_docs = vectordb.similarity_search(query, k=30)
+    context = " ".join([doc.page_content for doc in relevant_docs])
+    
+    # Generate response using OpenAI
+    response = openai.chat.completions.create(
+        model="gpt-4",
+        messages=[
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": f"Given the document: {context}\n\nGenerate a response to the query: {query}"}
+        ],
+        max_tokens=300,
+        temperature=0.7,
+    )
+    
+    return response.choices[0].message.content.strip()
+
+# Create Gradio interface
+demo = gr.Interface(
+    fn=process_query,
+    inputs=[
+        gr.Textbox(label="Enter your question", placeholder="Type your question here...")
+    ],
     outputs=[
-        gr.Textbox(label="Generated Response"),
-        gr.Textbox(label="Retrieved Documents")
+        gr.Textbox(label="Answer")
     ],
-    title="RAG-Based Question Answering System",
-    description="Enter a question and retrieve relevant documents with AI-generated response."
+    title="RAG-Powered Question Answering System",
+    description="Ask questions and get answers based on the embedded document knowledge.",
+    examples=[
+        ["What role does T-cell count play in severe human adenovirus type 55 (HAdV-55) infection?"],
+        ["In what school district is Governor John R. Rogers High School located?"],
+        ["Is there a functional neural correlate of individual differences in cardiovascular reactivity?"]
+    ]
 )
 
-# ✅ Launch the Gradio App
+# Launch the app
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()