import gradio as gr
from transformers import pipeline
from langchain_community.vectorstores import FAISS  # Updated import
from langchain_community.embeddings import HuggingFaceEmbeddings  # Updated import

# Load models
generator = pipeline("text-generation", model="gpt2")
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
vector_db = FAISS.load_local("faiss_index", embeddings)

def rag_response(question):
    # Retrieve relevant context
    docs = vector_db.similarity_search(question, k=2)
    context = "\n".join([d.page_content for d in docs])
    
    # Generate answer
    prompt = f"""Answer the question using this context:
{context}
Question: {question}
Answer:"""
    
    result = generator(
        prompt,
        max_length=200,
        num_return_sequences=1,
        temperature=0.3
    )
    
    return result[0]['generated_text']

# Create Gradio interface
iface = gr.Interface(
    fn=rag_response,
    inputs="text",
    outputs="text",
    title="RAG Demo",
    examples=["What is photosynthesis?", "Who discovered penicillin?"]
)

iface.launch()