Spaces:
Running
Running
import gradio as gr | |
import torch | |
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' | |
embedding_model = HuggingFaceEmbeddings(model_name=model_name) | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
# Load ChromaDB | |
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}" | |
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) | |
# Create Gradio App | |
iface = gr.Interface( | |
fn=gradio_interface, | |
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." | |
) | |
iface.launch() | |