| import torch | |
| import gradio as gr | |
| import faiss | |
| import numpy as np | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
| from sentence_transformers import SentenceTransformer | |
| GEN_MODEL_NAME = "google/flan-t5-small" | |
| EMB_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2" | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| tokenizer = AutoTokenizer.from_pretrained(GEN_MODEL_NAME) | |
| gen_model = AutoModelForSeq2SeqLM.from_pretrained(GEN_MODEL_NAME).to(DEVICE) | |
| embedder = SentenceTransformer(EMB_MODEL_NAME, device=DEVICE) | |
| DISCLAIMER = ( | |
| "Disclaimer: This pregnancy assistant provides general educational information only. " | |
| "It is not a substitute for antenatal care or professional medical advice. " | |
| "If you have warning signs (heavy bleeding, severe headache, severe abdominal pain, reduced baby movements) " | |
| "or feel unsafe, please go to a clinic or hospital immediately." | |
| ) | |
| pregnancy_facts = [ | |
| { | |
| "id": 1, | |
| "topic": "Antenatal care", | |
| "text": ( | |
| "Pregnant women should attend regular antenatal visits to monitor their own health " | |
| "and the baby's development. Early booking in the first trimester is recommended so " | |
| "that any health issues can be detected and managed early. During antenatal visits, " | |
| "healthcare providers may check blood pressure, weight, urine, and the baby's heartbeat, " | |
| "and may request simple laboratory tests or scans when needed. These visits are also a " | |
| "good time to ask questions about nutrition, exercise, birth preparation, and any symptoms " | |
| "that are worrying you. Skipping antenatal care can make it easier to miss important warning " | |
| "signs, so it is better to keep appointments and follow the advice of qualified health professionals." | |
| ) | |
| }, | |
| { | |
| "id": 2, | |
| "topic": "Warning signs", | |
| "text": ( | |
| "Some warning signs in pregnancy need urgent medical attention. These include severe or sudden " | |
| "headache, blurred vision, swelling of the face or hands, heavy vaginal bleeding, severe abdominal " | |
| "or pelvic pain, or fluid leaking from the vagina before labour starts. Other serious signs are fever, " | |
| "painful urination, convulsions, difficulty breathing, or a sudden reduction or change in the baby's " | |
| "movements. If a pregnant woman notices any of these symptoms, she should go to a clinic or hospital " | |
| "as soon as possible or call a healthcare professional for guidance. It is safer to be checked and told " | |
| "that everything is fine than to ignore a possible emergency." | |
| ) | |
| }, | |
| { | |
| "id": 3, | |
| "topic": "Nutrition", | |
| "text": ( | |
| "During pregnancy, eating a balanced diet helps support the health of the mother and the growing baby. " | |
| "Meals should include vegetables, fruits, whole grains, and sources of protein such as beans, eggs, fish, " | |
| "or lean meat, depending on what is culturally acceptable and available. Pregnant women are often advised " | |
| "by healthcare providers to take folic acid and iron supplements to help reduce the risk of anaemia and " | |
| "support the baby's development. Drinking enough clean water throughout the day is also important. " | |
| "Highly processed foods, very sugary drinks, or large amounts of junk food should be limited. " | |
| "For any supplements or herbal products, it is safer to ask a qualified healthcare professional before use." | |
| ) | |
| }, | |
| { | |
| "id": 4, | |
| "topic": "Substances to avoid", | |
| "text": ( | |
| "Alcohol, tobacco, and illicit or non-prescribed drugs should be avoided in pregnancy because they may " | |
| "harm the baby and affect the mother's health. Smoking or exposure to second-hand smoke can increase the " | |
| "risk of complications in pregnancy and after birth. Some medicines that are safe for non-pregnant adults " | |
| "may not be safe in pregnancy, especially if taken without medical advice. Before starting, stopping, or " | |
| "changing any medicine, including herbal remedies or over-the-counter drugs, a pregnant woman should speak " | |
| "to a qualified healthcare professional. It is better to check first than to take a product that could be " | |
| "unsafe for the baby." | |
| ) | |
| }, | |
| { | |
| "id": 5, | |
| "topic": "Exercise", | |
| "text": ( | |
| "For many women, moderate physical activity during pregnancy can be safe and beneficial. Simple activities " | |
| "like walking, gentle stretching, or pregnancy-safe exercises can help improve mood, sleep, and general " | |
| "fitness. However, very intense exercise, activities with a high risk of falling, or contact sports may not " | |
| "be suitable in pregnancy. Pregnant women who were not exercising before pregnancy should start slowly and " | |
| "listen to their bodies, stopping if they feel pain, dizziness, or shortness of breath. It is always a good " | |
| "idea to discuss exercise plans with a healthcare professional, especially if there are any medical problems " | |
| "or previous pregnancy complications." | |
| ) | |
| }, | |
| ] | |
| corpus = [item["text"] for item in pregnancy_facts] | |
| corpus_embeddings = embedder.encode(corpus, convert_to_numpy=True) | |
| dimension = corpus_embeddings.shape[1] | |
| index = faiss.IndexFlatL2(dimension) | |
| index.add(corpus_embeddings) | |
| def retrieve_pregnancy_facts(query: str, k: int = 3): | |
| query_embedding = embedder.encode([query], convert_to_numpy=True) | |
| distances, indices = index.search(query_embedding, k) | |
| results = [] | |
| for idx in indices[0]: | |
| results.append(pregnancy_facts[idx]["text"]) | |
| return results | |
| def pregnancy_qa_pipeline(question: str) -> str: | |
| question = question.strip() | |
| if not question: | |
| return "Please type a question about pregnancy." | |
| contexts = retrieve_pregnancy_facts(question, k=3) | |
| context_str = "\n\n".join(contexts) | |
| prompt = f"""You are an educational pregnancy information assistant. | |
| Context: | |
| {context_str} | |
| Instructions: | |
| - Give general educational information only. | |
| - Do not diagnose or confirm pregnancy status. | |
| - Do not provide medication doses, prescriptions, or step-by-step instructions for taking medicines. | |
| - Do not provide instructions for ending a pregnancy. | |
| - If the question sounds urgent or includes warning signs, strongly recommend immediate in-person care. | |
| - If the context is insufficient, say what you can safely say and recommend seeing a healthcare professional. | |
| Question: {question} | |
| Answer:""" | |
| inputs = tokenizer( | |
| prompt, | |
| return_tensors="pt", | |
| truncation=True, | |
| max_length=1024 | |
| ).to(DEVICE) | |
| with torch.no_grad(): | |
| outputs = gen_model.generate( | |
| **inputs, | |
| max_new_tokens=256, | |
| num_beams=4, | |
| early_stopping=True | |
| ) | |
| answer = tokenizer.decode(outputs[0], skip_special_tokens=True).strip() | |
| return f"{answer}\n\n{DISCLAIMER}" | |
| def pregnancy_qa_interface(question: str) -> str: | |
| return pregnancy_qa_pipeline(question) | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# 🤰 Pregnancy Information Assistant (Educational Use Only)") | |
| gr.Markdown( | |
| "Ask general questions about pregnancy, antenatal care, warning signs, and lifestyle. " | |
| "This tool is for education only and does not replace a healthcare professional." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| q_box = gr.Textbox( | |
| label="Ask a question about pregnancy", | |
| placeholder="e.g. Is it safe to exercise during pregnancy?", | |
| lines=2 | |
| ) | |
| ask_btn = gr.Button("Ask") | |
| with gr.Column(): | |
| a_box = gr.Textbox( | |
| label="Answer", | |
| lines=10 | |
| ) | |
| ask_btn.click( | |
| pregnancy_qa_interface, | |
| inputs=q_box, | |
| outputs=a_box | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |