8 changed files with 536 additions and 0 deletions
@ -0,0 +1,37 @@ |
|||||||
|
# Overview |
||||||
|
|
||||||
|
This uses de-identified medical dictation data supplied by [mtsamples](https://mtsamples.com). The data from the mtsamples |
||||||
|
website was download from [kaggle](https://www.kaggle.com/datasets/tboyle10/medicaltranscriptions). There are four |
||||||
|
sample notes in different directories (see knowledge_base/mtsamples_dictations) that will added to a chromaDb |
||||||
|
vector database and will be available during chat using RAG (Retrieval Augmented Generation). |
||||||
|
|
||||||
|
# How to run |
||||||
|
|
||||||
|
- Run example |
||||||
|
|
||||||
|
```shell |
||||||
|
conda activate <your_environment> |
||||||
|
cd <your_directory_where_script_lives> |
||||||
|
python run_rag_chat.py |
||||||
|
``` |
||||||
|
|
||||||
|
# Chat example |
||||||
|
|
||||||
|
 |
||||||
|
|
||||||
|
# Questions to ask? |
||||||
|
|
||||||
|
1) How old is Ms. Connor? |
||||||
|
2) What are Ms. Connor's vital signs? |
||||||
|
3) How old is Ms. Mouse? |
||||||
|
4) What is Ms. Mouse concerned about? |
||||||
|
5) What are Ms. Mouse's vital signs? |
||||||
|
6) How old is Mr. Duck? |
||||||
|
7) Why did Mr. Duck go to the doctor? |
||||||
|
8) How old is Ms. Barbara? |
||||||
|
9) Why did Ms. Barbara go to the doctor? |
||||||
|
10) Is Ms. Barbara allergic to anything? |
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
After Width: | Height: | Size: 89 KiB |
@ -0,0 +1,44 @@ |
|||||||
|
HISTORY OF PRESENT ILLNESS: |
||||||
|
|
||||||
|
Ms. Connor is a 50-year-old female who returns to clinic for a wound check. |
||||||
|
The patient underwent an APR secondary to refractory ulcerative colitis. |
||||||
|
Subsequently, she developed a wound infection, which has since healed. |
||||||
|
On our most recent visit to our clinic, she has her perineal stitches removed and presents today for followup of |
||||||
|
her perineal wound. She describes no drainage or erythema from her bottom. She is having good ostomy output. |
||||||
|
She does not describe any fevers, chills, nausea, or vomiting. The patient does describe some intermittent |
||||||
|
pain beneath the upper portion of the incision as well as in the right lower quadrant below her ostomy. |
||||||
|
She has been taking Percocet for this pain and it does work. She has since run out has been trying |
||||||
|
extra strength Tylenol, which will occasionally help this intermittent pain. She is requesting additional |
||||||
|
pain medications for this occasional abdominal pain, which she still experiences. |
||||||
|
|
||||||
|
PHYSICAL EXAMINATION: |
||||||
|
|
||||||
|
Temperature 95.8, pulse 68, blood pressure 132/73, and weight 159 pounds. |
||||||
|
|
||||||
|
This is a pleasant female in no acute distress. |
||||||
|
The patient's abdomen is soft, nontender, nondistended with a well-healed midline scar. |
||||||
|
There is an ileostomy in the right hemiabdomen, which is pink, patent, productive, and protuberant. |
||||||
|
There are no signs of masses or hernias over the patient's abdomen. |
||||||
|
|
||||||
|
ASSESSMENT AND PLAN: |
||||||
|
|
||||||
|
This is a pleasant 50-year-old female who has undergone an APR secondary to refractory ulcerative colitis. |
||||||
|
Overall, her quality of life has significantly improved since she had her APR. She is functioning well with her ileostomy. |
||||||
|
She did have concerns or questions about her diet and we discussed the BRAT diet, which consisted of foods that would |
||||||
|
slow down the digestive tract such as bananas, rice, toast, cheese, and peanut butter. |
||||||
|
I discussed the need to monitor her ileostomy output and preferential amount of daily output is 2 liters or less. |
||||||
|
I have counseled her on refraining from soft drinks and fruit drinks. I have also discussed with her that this diet |
||||||
|
is moreover a trial and error and that she may try certain foods that did not agree with her ileostomy, |
||||||
|
however others may and that this is something she will just have to perform trials with over the next several |
||||||
|
months until she finds what foods that she can and cannot eat with her ileostomy. She also had questions about |
||||||
|
her occasional abdominal pain. I told her that this was probably continue to improve as months went by and I |
||||||
|
gave her a refill of her Percocet for the continued occasional pain. I told her that this would the last time |
||||||
|
I would refill the Percocet and if she has continued pain after she finishes this bottle then she would need to |
||||||
|
start ibuprofen or Tylenol if she had continued pain. The patient then brought up some right hand and arm numbness, |
||||||
|
which has been there postsurgically and was thought to be from positioning during surgery. |
||||||
|
This is all primarily gone away except for a little bit of numbness at the tip of the third digit as well as |
||||||
|
some occasional forearm muscle cramping. I told her that I felt that this would continue to improve as it |
||||||
|
has done over the past two months since her surgery. I told her to continue doing hand exercises as she has |
||||||
|
been doing and this seems to be working for her. Overall, I think she has healed from her surgery and is doing |
||||||
|
very well. Again, her quality of life is significantly improved. She is happy with her performance. We will see |
||||||
|
her back in six months just for a general routine checkup and see how she is doing at that time. |
@ -0,0 +1,50 @@ |
|||||||
|
HISTORY OF PRESENT ILLNESS:, |
||||||
|
|
||||||
|
Ms. Mouse is a 67-year-old white female with a history of uterine papillary serous carcinoma who is |
||||||
|
status post 6 cycles of carboplatin and Taxol, is here today for followup. Her last cycle of chemotherapy |
||||||
|
was finished on 01/18/08, and she complains about some numbness in her right upper extremity. |
||||||
|
This has not gotten worse recently and there is no numbness in her toes. She denies any tingling or burning., |
||||||
|
|
||||||
|
REVIEW OF SYSTEMS: |
||||||
|
|
||||||
|
Negative for any fever, chills, nausea, vomiting, headache, chest pain, shortness of breath, abdominal pain, |
||||||
|
constipation, diarrhea, melena, hematochezia or dysuria. |
||||||
|
|
||||||
|
The patient is concerned about her blood pressure being up a little bit and also a mole that she had noticed for the |
||||||
|
past few months in her head. |
||||||
|
|
||||||
|
PHYSICAL EXAMINATION: |
||||||
|
|
||||||
|
VITAL SIGNS: Temperature 35.6, blood pressure 143/83, pulse 65, respirations 18, and weight 66.5 kg. |
||||||
|
GENERAL: She is a middle-aged white female, not in any distress. |
||||||
|
HEENT: No lymphadenopathy or mucositis. |
||||||
|
CARDIOVASCULAR: Regular rate and rhythm. |
||||||
|
LUNGS: Clear to auscultation bilaterally. |
||||||
|
EXTREMITIES: No cyanosis, clubbing or edema. |
||||||
|
NEUROLOGICAL: No focal deficits noted. |
||||||
|
PELVIC: Normal-appearing external genitalia. Vaginal vault with no masses or bleeding., |
||||||
|
|
||||||
|
LABORATORY DATA: |
||||||
|
|
||||||
|
None today. |
||||||
|
|
||||||
|
RADIOLOGIC DATA: |
||||||
|
|
||||||
|
CT of the chest, abdomen, and pelvis from 01/28/08 revealed status post total abdominal hysterectomy/bilateral |
||||||
|
salpingo-oophorectomy with an unremarkable vaginal cuff. No local or distant metastasis. |
||||||
|
Right probably chronic gonadal vein thrombosis. |
||||||
|
|
||||||
|
ASSESSMENT: |
||||||
|
|
||||||
|
This is a 67-year-old white female with history of uterine papillary serous carcinoma, status post total |
||||||
|
abdominal hysterectomy and bilateral salpingo-oophorectomy and 6 cycles of carboplatin and Taxol chemotherapy. |
||||||
|
She is doing well with no evidence of disease clinically or radiologically. |
||||||
|
|
||||||
|
PLAN: |
||||||
|
|
||||||
|
1. Plan to follow her every 3 months and CT scans every 6 months for the first 2 years. |
||||||
|
2. The patient was advised to contact the primary physician for repeat blood pressure check and get started on |
||||||
|
antihypertensives if it is persistently elevated. |
||||||
|
3. The patient was told that the mole that she is mentioning in her head is no longer palpable and just to observe it for now. |
||||||
|
4. The patient was advised about doing Kegel exercises for urinary incontinence, and we will address this issue again |
||||||
|
during next clinic visit if it is persistent. |
@ -0,0 +1,25 @@ |
|||||||
|
SUBJECTIVE: |
||||||
|
|
||||||
|
Mr. Duck is a 29-year-old white male who is a patient of Dr. XYZ and he comes in today |
||||||
|
complaining that he was stung by a Yellow Jacket Wasp yesterday and now has a lot of |
||||||
|
swelling in his right hand and right arm. He says that he has been stung by wasps before and had similar |
||||||
|
reactions. He just said that he wanted to catch it early before he has too bad of a severe reaction like he has had in the past. |
||||||
|
He has had a lot of swelling, but no anaphylaxis-type reactions in the past; no shortness of breath or difficultly with his |
||||||
|
throat feeling like it is going to close up or anything like that in the past; no racing heart beat or anxiety feeling, |
||||||
|
just a lot of localized swelling where the sting occurs. |
||||||
|
|
||||||
|
OBJECTIVE: |
||||||
|
|
||||||
|
Vitals: His temperature is 98.4. Respiratory rate is 18. Weight is 250 pounds. |
||||||
|
Extremities: Examination of his right hand and forearm reveals that he has an apparent sting just around his |
||||||
|
wrist region on his right hand on the medial side as well as significant swelling in his hand and his right forearm; |
||||||
|
extending up to the elbow. He says that it is really not painful or anything like that. It is really not all that |
||||||
|
red and no signs of infection at this time. |
||||||
|
|
||||||
|
ASSESSMENT:, Wasp sting to the right wrist area. |
||||||
|
|
||||||
|
PLAN: |
||||||
|
|
||||||
|
1. Solu-Medrol 125 mg IM X 1. |
||||||
|
2. Over-the-counter Benadryl, ice and elevation of that extremity. |
||||||
|
3. Follow up with Dr. XYZ if any further evaluation is needed. |
@ -0,0 +1,54 @@ |
|||||||
|
CHIEF COMPLAINT: |
||||||
|
|
||||||
|
Ms. Barbara is a thirty one year old female patient comes for three-week postpartum checkup, complaining of allergies. |
||||||
|
|
||||||
|
HISTORY OF PRESENT ILLNESS: |
||||||
|
|
||||||
|
She is doing well postpartum. She has had no headache. She is breastfeeding and feels like her milk is adequate. |
||||||
|
She has not had much bleeding. She is using about a mini pad twice a day, not any cramping or clotting and the |
||||||
|
discharge is turned from red to brown to now slightly yellowish. She has not yet had sexual intercourse. |
||||||
|
She does complain that she has had a little pain with the bowel movement, and every now and then she |
||||||
|
notices a little bright red bleeding. She has not been particularly constipated but her husband says |
||||||
|
she is not eating her vegetables like she should. Her seasonal allergies have back developed and she is |
||||||
|
complaining of extremely itchy watery eyes, runny nose, sneezing, and kind of a pressure sensation in her ears. |
||||||
|
|
||||||
|
MEDICATIONS: |
||||||
|
|
||||||
|
Prenatal vitamins. |
||||||
|
|
||||||
|
ALLERGIES: |
||||||
|
|
||||||
|
She thinks to Benadryl. |
||||||
|
|
||||||
|
FAMILY HISTORY: |
||||||
|
|
||||||
|
Mother is 50 and healthy. Dad is 40 and healthy. Half-sister, age 34, is healthy. |
||||||
|
She has a sister who is age 10 who has some yeast infections. |
||||||
|
|
||||||
|
PHYSICAL EXAMINATION: |
||||||
|
|
||||||
|
VITALS: Weight: 124 pounds. Blood pressure 96/54. Pulse: 72. Respirations: 16. LMP: 10/18/03. Age: 39. |
||||||
|
HEENT: Head is normocephalic. |
||||||
|
Eyes: EOMs intact. |
||||||
|
PERRLA. Conjunctiva clear. |
||||||
|
Fundi: Discs flat, cups normal. |
||||||
|
No AV nicking, hemorrhage or exudate. |
||||||
|
Ears: TMs intact. |
||||||
|
Mouth: No lesion. |
||||||
|
Throat: No inflammation. |
||||||
|
She has allergic rhinitis with clear nasal drainage, clear watery discharge from the eyes. |
||||||
|
Abdomen: Soft. No masses. |
||||||
|
Pelvic: Uterus is involuting. |
||||||
|
Rectal: She has one external hemorrhoid which has inflamed. Stool is guaiac negative and using anoscope, |
||||||
|
no other lesions are identified. |
||||||
|
|
||||||
|
ASSESSMENT/PLAN: |
||||||
|
|
||||||
|
Satisfactory three-week postpartum course, seasonal allergies. We will try Patanol eyedrops and Allegra 60 |
||||||
|
mg twice a day. She was cautioned about the possibility that this may alter her milk supply. She is to |
||||||
|
drink extra fluids and call if she has problems with that. We will try ProctoFoam HC. For the hemorrhoids, |
||||||
|
also increase the fiber in her diet. That prescription was written, as well as one for Allegra and Patanol. |
||||||
|
She additionally will be begin on Micronor because she would like to protect herself from pregnancy until |
||||||
|
her husband get scheduled in and has a vasectomy, which is their ultimate plan for birth control, and she |
||||||
|
anticipates that happening fairly soon. She will call and return if she continues to have problems with allergies. |
||||||
|
Meantime, rechecking in three weeks for her final six-week postpartum checkup. |
@ -0,0 +1,59 @@ |
|||||||
|
import gradio as gr |
||||||
|
from langchain_chroma import Chroma |
||||||
|
from pathlib import Path |
||||||
|
from utils import create_vector_db, Rag, get_chunks, get_conversation_chain, get_local_vector_db |
||||||
|
|
||||||
|
|
||||||
|
def chat(question, history) -> str: |
||||||
|
|
||||||
|
""" |
||||||
|
Get the chat data need for the gradio app |
||||||
|
|
||||||
|
:param question: |
||||||
|
The question being asked in the chat app. |
||||||
|
:type question: str |
||||||
|
:param history: |
||||||
|
A list of the conversation questions and answers. |
||||||
|
:type history: list |
||||||
|
:return: |
||||||
|
The answer from the current question. |
||||||
|
""" |
||||||
|
|
||||||
|
result = conversation_chain.invoke({"question": question}) |
||||||
|
answer = result['answer'] |
||||||
|
|
||||||
|
# include source documents if they exist |
||||||
|
# grab the first one as that should be related to the answer |
||||||
|
source_doc = "" |
||||||
|
if result.get('source_documents'): |
||||||
|
source_doc = result['source_documents'][0] |
||||||
|
|
||||||
|
response = f"{answer}\n\n**Source:**\n{source_doc.metadata.get('source', 'Source')}" \ |
||||||
|
if source_doc \ |
||||||
|
else answer |
||||||
|
return response |
||||||
|
|
||||||
|
|
||||||
|
def main(): |
||||||
|
|
||||||
|
gr.ChatInterface(chat, type="messages").launch(inbrowser=True) |
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__': |
||||||
|
|
||||||
|
create_new_db = False if Path('vector_db').exists() else True |
||||||
|
|
||||||
|
if create_new_db: |
||||||
|
folders = Path('knowledge_base').glob('*') |
||||||
|
chunks = get_chunks(folders=folders) |
||||||
|
vector_store = create_vector_db(chunks=chunks, db_name=Rag.DB_NAME.value, embeddings=Rag.EMBED_MODEL.value) |
||||||
|
else: |
||||||
|
client = get_local_vector_db(path='../rag_chat_example/vector_db') |
||||||
|
vector_store = Chroma(client=client, embedding_function=Rag.EMBED_MODEL.value) |
||||||
|
|
||||||
|
conversation_chain = get_conversation_chain(vectorstore=vector_store) |
||||||
|
|
||||||
|
main() |
||||||
|
|
||||||
|
|
||||||
|
|
@ -0,0 +1,267 @@ |
|||||||
|
from chromadb import PersistentClient |
||||||
|
from dotenv import load_dotenv |
||||||
|
from enum import Enum |
||||||
|
|
||||||
|
import plotly.graph_objects as go |
||||||
|
from langchain.document_loaders import DirectoryLoader, TextLoader |
||||||
|
from langchain.text_splitter import CharacterTextSplitter |
||||||
|
from langchain.schema import Document |
||||||
|
from langchain_openai import OpenAIEmbeddings, ChatOpenAI |
||||||
|
from langchain_chroma import Chroma |
||||||
|
from langchain.memory import ConversationBufferMemory |
||||||
|
from langchain.chains import ConversationalRetrievalChain |
||||||
|
import numpy as np |
||||||
|
import os |
||||||
|
from pathlib import Path |
||||||
|
from sklearn.manifold import TSNE |
||||||
|
from typing import Any, List, Tuple, Generator |
||||||
|
|
||||||
|
cur_path = Path(__file__) |
||||||
|
env_path = cur_path.parent.parent.parent.parent / '.env' |
||||||
|
assert env_path.exists(), f"Please add an .env to the root project path" |
||||||
|
|
||||||
|
load_dotenv(dotenv_path=env_path) |
||||||
|
|
||||||
|
|
||||||
|
class Rag(Enum): |
||||||
|
|
||||||
|
GPT_MODEL = "gpt-4o-mini" |
||||||
|
HUG_MODEL = "sentence-transformers/all-MiniLM-L6-v2" |
||||||
|
EMBED_MODEL = OpenAIEmbeddings() |
||||||
|
DB_NAME = "vector_db" |
||||||
|
|
||||||
|
|
||||||
|
def add_metadata(doc: Document, doc_type: str) -> Document: |
||||||
|
""" |
||||||
|
Add metadata to a Document object. |
||||||
|
|
||||||
|
:param doc: The Document object to add metadata to. |
||||||
|
:type doc: Document |
||||||
|
:param doc_type: The type of document to be added as metadata. |
||||||
|
:type doc_type: str |
||||||
|
:return: The Document object with added metadata. |
||||||
|
:rtype: Document |
||||||
|
""" |
||||||
|
doc.metadata["doc_type"] = doc_type |
||||||
|
return doc |
||||||
|
|
||||||
|
|
||||||
|
def get_chunks(folders: Generator[Path, None, None], file_ext='.txt') -> List[Document]: |
||||||
|
""" |
||||||
|
Load documents from specified folders, add metadata, and split them into chunks. |
||||||
|
|
||||||
|
:param folders: List of folder paths containing documents. |
||||||
|
:type folders: List[str] |
||||||
|
:param file_ext: |
||||||
|
The file extension to get from a local knowledge base (e.g. '.txt') |
||||||
|
:type file_ext: str |
||||||
|
:return: List of document chunks. |
||||||
|
:rtype: List[Document] |
||||||
|
""" |
||||||
|
text_loader_kwargs = {'encoding': 'utf-8'} |
||||||
|
documents = [] |
||||||
|
for folder in folders: |
||||||
|
doc_type = os.path.basename(folder) |
||||||
|
loader = DirectoryLoader( |
||||||
|
folder, glob=f"**/*{file_ext}", loader_cls=TextLoader, loader_kwargs=text_loader_kwargs |
||||||
|
) |
||||||
|
folder_docs = loader.load() |
||||||
|
documents.extend([add_metadata(doc, doc_type) for doc in folder_docs]) |
||||||
|
|
||||||
|
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200) |
||||||
|
chunks = text_splitter.split_documents(documents) |
||||||
|
|
||||||
|
return chunks |
||||||
|
|
||||||
|
|
||||||
|
def create_vector_db(db_name: str, chunks: List[Document], embeddings: Any) -> Any: |
||||||
|
""" |
||||||
|
Create a vector database from document chunks. |
||||||
|
|
||||||
|
:param db_name: Name of the database to create. |
||||||
|
:type db_name: str |
||||||
|
:param chunks: List of document chunks. |
||||||
|
:type chunks: List[Document] |
||||||
|
:param embeddings: Embedding function to use. |
||||||
|
:type embeddings: Any |
||||||
|
:return: Created vector store. |
||||||
|
:rtype: Any |
||||||
|
""" |
||||||
|
# Delete if already exists |
||||||
|
if os.path.exists(db_name): |
||||||
|
Chroma(persist_directory=db_name, embedding_function=embeddings).delete_collection() |
||||||
|
|
||||||
|
# Create vectorstore |
||||||
|
vectorstore = Chroma.from_documents(documents=chunks, embedding=embeddings, persist_directory=db_name) |
||||||
|
|
||||||
|
return vectorstore |
||||||
|
|
||||||
|
|
||||||
|
def get_local_vector_db(path: str) -> Any: |
||||||
|
""" |
||||||
|
Get a local vector database. |
||||||
|
|
||||||
|
:param path: Path to the local vector database. |
||||||
|
:type path: str |
||||||
|
:return: Persistent client for the vector database. |
||||||
|
:rtype: Any |
||||||
|
""" |
||||||
|
return PersistentClient(path=path) |
||||||
|
|
||||||
|
|
||||||
|
def get_vector_db_info(vector_store: Any) -> None: |
||||||
|
""" |
||||||
|
Print information about the vector database. |
||||||
|
|
||||||
|
:param vector_store: Vector store to get information from. |
||||||
|
:type vector_store: Any |
||||||
|
""" |
||||||
|
collection = vector_store._collection |
||||||
|
count = collection.count() |
||||||
|
|
||||||
|
sample_embedding = collection.get(limit=1, include=["embeddings"])["embeddings"][0] |
||||||
|
dimensions = len(sample_embedding) |
||||||
|
|
||||||
|
print(f"There are {count:,} vectors with {dimensions:,} dimensions in the vector store") |
||||||
|
|
||||||
|
|
||||||
|
def get_plot_data(collection: Any) -> Tuple[np.ndarray, List[str], List[str], List[str]]: |
||||||
|
""" |
||||||
|
Get plot data from a collection. |
||||||
|
|
||||||
|
:param collection: Collection to get data from. |
||||||
|
:type collection: Any |
||||||
|
:return: Tuple containing vectors, colors, document types, and documents. |
||||||
|
:rtype: Tuple[np.ndarray, List[str], List[str], List[str]] |
||||||
|
""" |
||||||
|
result = collection.get(include=['embeddings', 'documents', 'metadatas']) |
||||||
|
vectors = np.array(result['embeddings']) |
||||||
|
documents = result['documents'] |
||||||
|
metadatas = result['metadatas'] |
||||||
|
doc_types = [metadata['doc_type'] for metadata in metadatas] |
||||||
|
colors = [['blue', 'green', 'red', 'orange'][['products', 'employees', 'contracts', 'company'].index(t)] for t in |
||||||
|
doc_types] |
||||||
|
|
||||||
|
return vectors, colors, doc_types, documents |
||||||
|
|
||||||
|
|
||||||
|
def get_2d_plot(collection: Any) -> go.Figure: |
||||||
|
""" |
||||||
|
Generate a 2D plot of the vector store. |
||||||
|
|
||||||
|
:param collection: Collection to generate plot from. |
||||||
|
:type collection: Any |
||||||
|
:return: 2D scatter plot figure. |
||||||
|
:rtype: go.Figure |
||||||
|
""" |
||||||
|
vectors, colors, doc_types, documents = get_plot_data(collection) |
||||||
|
tsne = TSNE(n_components=2, random_state=42) |
||||||
|
reduced_vectors = tsne.fit_transform(vectors) |
||||||
|
|
||||||
|
fig = go.Figure(data=[go.Scatter( |
||||||
|
x=reduced_vectors[:, 0], |
||||||
|
y=reduced_vectors[:, 1], |
||||||
|
mode='markers', |
||||||
|
marker=dict(size=5, color=colors, opacity=0.8), |
||||||
|
text=[f"Type: {t}<br>Text: {d[:100]}..." for t, d in zip(doc_types, documents)], |
||||||
|
hoverinfo='text' |
||||||
|
)]) |
||||||
|
|
||||||
|
fig.update_layout( |
||||||
|
title='2D Chroma Vector Store Visualization', |
||||||
|
scene=dict(xaxis_title='x', yaxis_title='y'), |
||||||
|
width=800, |
||||||
|
height=600, |
||||||
|
margin=dict(r=20, b=10, l=10, t=40) |
||||||
|
) |
||||||
|
|
||||||
|
return fig |
||||||
|
|
||||||
|
|
||||||
|
def get_3d_plot(collection: Any) -> go.Figure: |
||||||
|
""" |
||||||
|
Generate a 3D plot of the vector store. |
||||||
|
|
||||||
|
:param collection: Collection to generate plot from. |
||||||
|
:type collection: Any |
||||||
|
:return: 3D scatter plot figure. |
||||||
|
:rtype: go.Figure |
||||||
|
""" |
||||||
|
vectors, colors, doc_types, documents = get_plot_data(collection) |
||||||
|
tsne = TSNE(n_components=3, random_state=42) |
||||||
|
reduced_vectors = tsne.fit_transform(vectors) |
||||||
|
|
||||||
|
fig = go.Figure(data=[go.Scatter3d( |
||||||
|
x=reduced_vectors[:, 0], |
||||||
|
y=reduced_vectors[:, 1], |
||||||
|
z=reduced_vectors[:, 2], |
||||||
|
mode='markers', |
||||||
|
marker=dict(size=5, color=colors, opacity=0.8), |
||||||
|
text=[f"Type: {t}<br>Text: {d[:100]}..." for t, d in zip(doc_types, documents)], |
||||||
|
hoverinfo='text' |
||||||
|
)]) |
||||||
|
|
||||||
|
fig.update_layout( |
||||||
|
title='3D Chroma Vector Store Visualization', |
||||||
|
scene=dict(xaxis_title='x', yaxis_title='y', zaxis_title='z'), |
||||||
|
width=900, |
||||||
|
height=700, |
||||||
|
margin=dict(r=20, b=10, l=10, t=40) |
||||||
|
) |
||||||
|
|
||||||
|
return fig |
||||||
|
|
||||||
|
|
||||||
|
def get_conversation_chain(vectorstore: Any) -> ConversationalRetrievalChain: |
||||||
|
""" |
||||||
|
Create a conversation chain using the vector store. |
||||||
|
|
||||||
|
:param vectorstore: Vector store to use in the conversation chain. |
||||||
|
:type vectorstore: Any |
||||||
|
:return: Conversational retrieval chain. |
||||||
|
:rtype: ConversationalRetrievalChain |
||||||
|
""" |
||||||
|
llm = ChatOpenAI(temperature=0.7, model_name=Rag.GPT_MODEL.value) |
||||||
|
|
||||||
|
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True, output_key='answer') |
||||||
|
|
||||||
|
retriever = vectorstore.as_retriever(search_kwargs={"k": 25}) |
||||||
|
|
||||||
|
conversation_chain = ConversationalRetrievalChain.from_llm( |
||||||
|
llm=llm, |
||||||
|
retriever=retriever, |
||||||
|
memory=memory, |
||||||
|
return_source_documents=True, |
||||||
|
) |
||||||
|
|
||||||
|
return conversation_chain |
||||||
|
|
||||||
|
|
||||||
|
def get_lang_doc(document_text, doc_id, metadata=None, encoding='utf-8'): |
||||||
|
|
||||||
|
""" |
||||||
|
Build a langchain Document that can be used to create a chroma database |
||||||
|
|
||||||
|
:type document_text: str |
||||||
|
:param document_text: |
||||||
|
The text to add to a document object |
||||||
|
:type doc_id: str |
||||||
|
:param doc_id: |
||||||
|
The document id to include. |
||||||
|
:type metadata: dict |
||||||
|
:param metadata: |
||||||
|
A dictionary of metadata to associate to the document object. This will help filter an item from a |
||||||
|
vector database. |
||||||
|
:type encoding: string |
||||||
|
:param encoding: |
||||||
|
The type of encoding to use for loading the text. |
||||||
|
|
||||||
|
""" |
||||||
|
return Document( |
||||||
|
page_content=document_text, |
||||||
|
id=doc_id, |
||||||
|
metadata=metadata, |
||||||
|
encoding=encoding, |
||||||
|
) |
||||||
|
|
||||||
|
|
Loading…
Reference in new issue