import gradio as gr
from sentence_transformers import SentenceTransformer, util
from transformers import pipeline
 
# this loads the documents
with open("info.txt", "r") as f:
    docs = [line.strip() for line in f.readlines() if line.strip()]
 
# this loads the model for encoding 
# Encodes sentences into vectors so we can find the most relevant doc
embedder = SentenceTransformer("all-MiniLM-L6-v2")
doc_embeddings = embedder.encode(docs, convert_to_tensor=True)
 
# text generaition model I had found
qa_pipeline = pipeline("text-generation", model="gpt2")
 
# -Rag functionality
def answer_question(user_question):
    if not user_question.strip():
        return "Please type a question."
 
    # retrieves and finds the most relevant document line
    question_embedding = embedder.encode(user_question, convert_to_tensor=True)
    scores = util.cos_sim(question_embedding, doc_embeddings)[0]
    best_index = scores.argmax().item()
    retrieved_doc = docs[best_index]
 
    # builds a prompt and ask the model
    prompt = (
        f"Information: {retrieved_doc}\n\n"
        f"Question: {user_question}\n\n"
        f"Answer:"
    )
 
    output = qa_pipeline(prompt, max_new_tokens=150, do_sample=False)[0]["generated_text"]
 
    # Strip the prompt from the output so only the answer shows
    answer = output[len(prompt):].strip()
 
    return f"Answer:\n{answer}\n\n---\nSource used:\n{retrieved_doc}"
 
# Gradio UI 
demo = gr.Interface(
    fn=answer_question,
    inputs=gr.Textbox(
        lines=2,
        placeholder="e.g. How does racial bias affect Alzheimer's research?",
        label="Your question"
    ),
    outputs=gr.Textbox(label="Response", lines=10),
    title="AI Bias in Neurodegeneration Research — RAG Chatbot",
    description=(
        "Ask questions about racial bias in neurodegenerative disease research. "
        "This chatbot retrieves relevant information from a curated document set "
        "and generates a grounded answer."
    ),
    examples=[
        ["How does racial bias affect Alzheimer's research?"],
        ["Are Black patients underrepresented in Parkinson's clinical trials?"],
        ["What can AI do to reduce bias in neurodegeneration studies?"],
    ]
)
 
demo.launch()