File size: 3,301 Bytes
d6b5ead a12afd1 d6b5ead cff6007 d6b5ead cff6007 d6b5ead cff6007 d6b5ead cff6007 d6b5ead cff6007 d6b5ead cff6007 d6b5ead cff6007 d6b5ead cff6007 d6b5ead 35fa0eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import streamlit as st
from transformers import pipeline
import faiss
import numpy as np
import os
# Initialize models and retrieval pipeline
@st.cache_resource
def load_model():
# Load the retrieval and generation models
retrieval_model = pipeline('question-answering', model="bert-large-uncased-whole-word-masking-finetuned-squad")
generation_model = pipeline('text-generation', model="gpt-2")
return retrieval_model, generation_model
retrieval_model, generation_model = load_model()
# Function to search relevant medical documents using FAISS
def search_medical_documents(query, faiss_index):
# Convert the query to a vector (embedding) using a model (e.g., BERT or any embedding model)
query_vector = retrieval_model.tokenizer(query, return_tensors="pt", padding=True, truncation=True)
query_embedding = retrieval_model.model(**query_vector).last_hidden_state.mean(dim=1).detach().numpy()
# Perform search in FAISS index
D, I = faiss_index.search(query_embedding, k=5) # Top 5 results
retrieved_docs = [faiss_index.reconstruct(i) for i in I[0]] # Retrieve documents based on indices
return retrieved_docs
# Build FAISS index (you should have your own documents here)
def build_faiss_index(documents):
embeddings = [retrieval_model.tokenizer(doc, return_tensors="pt", padding=True, truncation=True) for doc in documents]
embeddings = [retrieval_model.model(**embed).last_hidden_state.mean(dim=1).detach().numpy() for embed in embeddings]
dim = len(embeddings[0]) # Dimensionality of the embedding
faiss_index = faiss.IndexFlatL2(dim)
faiss_index.add(np.array(embeddings)) # Add embeddings to the index
return faiss_index
# Set up the Streamlit app interface
def chatbot_interface():
st.title("AI-Powered Healthcare Chatbot")
st.write("Ask any health-related questions or upload medical records for analysis.")
user_input = st.text_input("Enter your query here:")
if user_input:
# Retrieve relevant documents using the FAISS index (after indexing documents)
retrieved_docs = search_medical_documents(user_input, faiss_index)
context = " ".join(retrieved_docs)
response = generation_model(f"Context: {context} \nQuery: {user_input}\nResponse:")
st.write("Response: ")
st.write(response[0]['generated_text'])
# File uploader for medical records
uploaded_file = st.file_uploader("Upload medical documents (PDF, DOCX, etc.):", type=["pdf", "docx"])
if uploaded_file is not None:
extracted_text = extract_text_from_file(uploaded_file)
st.write("Extracted Text: ")
st.write(extracted_text)
response = generation_model(f"Extracted Text: {extracted_text}\nAnswer the following question: {user_input}")
st.write("Generated Response:")
st.write(response[0]['generated_text'])
# Function to extract text from uploaded medical files
def extract_text_from_file(uploaded_file):
file_name = uploaded_file.name
if file_name.endswith(".pdf"):
return extract_text_from_pdf(uploaded_file)
elif file_name.endswith(".docx"):
return extract_text_from_docx(uploaded_file)
else:
return "File type not supported."
if __name__ == "__main__":
chatbot_interface() |