Spaces:

HFUsman
/

Doctor_AI

Build error

Update app.py

35fa0eb verified almost 2 years ago

3.3 kB

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()

	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()