import nltk from nltk.tokenize import word_tokenize from langchain_community.document_loaders import TextLoader from langchain_community.embeddings.sentence_transformer import ( SentenceTransformerEmbeddings, ) from langchain_community.vectorstores import Chroma from langchain_text_splitters import CharacterTextSplitter # Download NLTK data for tokenization nltk.download('punkt') import os global db class QuestionRetriever: def load_documents(self,file_name): data_directory = "data/" file_path = os.path.join(data_directory, file_name) loader = TextLoader(file_path) documents = loader.load() return documents def store_data_in_vector_db(self,documents): # global db text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=0,separator="\n") docs = text_splitter.split_documents(documents) # create the open-source embedding function embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2") # print(docs) # load it into Chroma db = Chroma.from_documents(docs, embedding_function) return db def get_response(self, user_query, predicted_mental_category): if predicted_mental_category == "depression": documents=self.load_documents("depression_questions.txt") elif predicted_mental_category == "adhd": documents=self.load_documents("adhd_questions.txt") elif predicted_mental_category == "anxiety": documents=self.load_documents("anxiety_questions.txt") else: print("Sorry, allowed predicted_mental_category is ['depresison', 'adhd', 'anxiety'].") return db=self.store_data_in_vector_db(documents) docs = db.similarity_search(user_query) most_similar_question = docs[0].page_content.split("\n")[0] # Extract the first question if user_query==most_similar_question: most_similar_question=docs[1].page_content.split("\n")[0] print(most_similar_question) return most_similar_question if __name__ == "__main__": model = QuestionRetriever() user_input = input("User: ") predicted_mental_condition = "depression" response = model.get_response(user_input, predicted_mental_condition) print("Chatbot:", response)