Model export changed to safetensors. Improved representational model function. Got zero shot topic modelling working

.gitignore

@@ -4,6 +4,10 @@
 *.csv
 *.pkl
 *.parquet
+*.png
+*.safetensors
+*.json
 .ipynb_checkpoints/*
 old_code/*
-model/*
+model/*
+output_model/*

app.py

@@ -1,17 +1,14 @@
-import os
-
-#os.environ["TOKENIZERS_PARALLELISM"] = "true"
-#os.environ["HF_HOME"] = "/mnt/c/..."
-#os.environ["CUDA_PATH"] = "/mnt/c/..."
-#print(os.environ["HF_HOME"])
-
 import gradio as gr
 from datetime import datetime
 import pandas as pd
 import numpy as np
 from sklearn.cluster import KMeans
 from sklearn.feature_extraction.text import CountVectorizer
-from transformers import AutoModel
+from transformers import AutoModel, AutoTokenizer
+from transformers.pipelines import pipeline
+from sklearn.pipeline import make_pipeline
+from sklearn.decomposition import TruncatedSVD
+from sklearn.feature_extraction.text import TfidfVectorizer
 import funcs.anonymiser as anon
 
 from torch import cuda, backends, version
@@ -43,10 +40,11 @@ from bertopic import BERTopic
 today = datetime.now().strftime("%d%m%Y")
 today_rev = datetime.now().strftime("%Y%m%d")
 
-from funcs.helper_functions import dummy_function, put_columns_in_df, read_file, get_file_path_end
-from funcs.representation_model import representation_model
+from funcs.helper_functions import dummy_function, put_columns_in_df, read_file, get_file_path_end, zip_folder, delete_files_in_folder
+#from funcs.representation_model import representation_model
 from funcs.embeddings import make_or_load_embeddings
 
+
 # Load embeddings
 #embedding_model_name = "BAAI/bge-small-en-v1.5"
 #embedding_model = SentenceTransformer(embedding_model_name)
@@ -57,14 +55,24 @@ embeddings_name = "jinaai/jina-embeddings-v2-small-en"
 local_embeddings_location = "model/jina/"
 revision_choice = "b811f03af3d4d7ea72a7c25c802b21fc675a5d99"
 
-try:
-    embedding_model = AutoModel.from_pretrained(local_embeddings_location, revision = revision_choice, trust_remote_code=True,local_files_only=True, device_map="auto")
-except:
-    embedding_model = AutoModel.from_pretrained(embeddings_name, revision = revision_choice, trust_remote_code=True, device_map="auto")
+if low_resource_mode == "No":
+    try:
+        embedding_model = AutoModel.from_pretrained(local_embeddings_location, revision = revision_choice, trust_remote_code=True,local_files_only=True, device_map="auto")
+    except:
+        embedding_model = AutoModel.from_pretrained(embeddings_name, revision = revision_choice, trust_remote_code=True, device_map="auto")
+
+    tokenizer = AutoTokenizer.from_pretrained("jinaai/jina-embeddings-v2-small-en")
+
+    embedding_model_pipe = pipeline("feature-extraction", model=embedding_model, tokenizer=tokenizer)
 
+elif low_resource_mode == "Yes":
+    embedding_model_pipe = make_pipeline(
+                TfidfVectorizer(),
+                TruncatedSVD(2) # 100 # set to 2 to be compatible with zero shot topics - can't be higher than number of topics
+                )
 
 
-def extract_topics(in_files, in_file, min_docs_slider, in_colnames, max_topics_slider, candidate_topics, in_label, anonymise_drop, return_intermediate_files, embeddings_super_compress, low_resource_mode_opt):
+def extract_topics(in_files, in_file, min_docs_slider, in_colnames, max_topics_slider, candidate_topics, in_label, anonymise_drop, return_intermediate_files, embeddings_super_compress, low_resource_mode, create_llm_topic_labels):
     
     file_list = [string.name for string in in_file]
 
@@ -91,7 +99,28 @@ def extract_topics(in_files, in_file, min_docs_slider, in_colnames, max_topics_s
     ## Load in pre-embedded file if exists
     file_list = [string.name for string in in_file]
 
-    embeddings_out, reduced_embeddings = make_or_load_embeddings(docs, file_list, data_file_name_no_ext, embedding_model, return_intermediate_files, embeddings_super_compress, low_resource_mode_opt)
+    print("Low resource mode: ", low_resource_mode)
+
+    if low_resource_mode == "No":
+        print("Choosing high resource Jina transformer model")
+        try:
+            embedding_model = AutoModel.from_pretrained(local_embeddings_location, revision = revision_choice, trust_remote_code=True,local_files_only=True, device_map="auto")
+        except:
+            embedding_model = AutoModel.from_pretrained(embeddings_name, revision = revision_choice, trust_remote_code=True, device_map="auto")
+
+        tokenizer = AutoTokenizer.from_pretrained("jinaai/jina-embeddings-v2-small-en")
+
+        embedding_model_pipe = pipeline("feature-extraction", model=embedding_model, tokenizer=tokenizer)
+
+    elif low_resource_mode == "Yes":
+        print("Choosing low resource TfIDF model")
+        embedding_model_pipe = make_pipeline(
+                TfidfVectorizer(),
+                TruncatedSVD(2) # 100 # To be compatible with zero shot, this needs to be lower than number of suggested topics
+                )
+        embedding_model = embedding_model_pipe
+
+    embeddings_out, reduced_embeddings = make_or_load_embeddings(docs, file_list, data_file_name_no_ext, embedding_model, return_intermediate_files, embeddings_super_compress, low_resource_mode, create_llm_topic_labels)
 
     # all_lengths = [len(embedding) for embedding in embeddings_out]
     # if len(set(all_lengths)) > 1:
@@ -110,18 +139,24 @@ def extract_topics(in_files, in_file, min_docs_slider, in_colnames, max_topics_s
     #cluster_model = KMeans(n_clusters=max_topics_slider)
 
     # Countvectoriser removes stopwords, combines terms up to 2 together:
-    if min_docs_slider < 3:
-        min_df_val = min_docs_slider
-    else:
-        min_df_val = 3
+    #if min_docs_slider < 3:
+    #    min_df_val = min_docs_slider
+    #else:
+    #    min_df_val = 3
 
-    print(min_df_val)
+    #print(min_df_val)
 
     vectoriser_model = CountVectorizer(stop_words="english", ngram_range=(1, 2), min_df=0.1)
 
+    
+    from funcs.prompts import capybara_prompt, capybara_start, open_hermes_prompt, open_hermes_start, stablelm_prompt, stablelm_start
+    from funcs.representation_model import create_representation_model, found_file, gpu_config, chosen_start_tag
+
+    print("Create LLM topic labels:", create_llm_topic_labels)
+    representation_model = create_representation_model(create_llm_topic_labels, gpu_config, found_file, chosen_start_tag)
 
     if not candidate_topics:
-        topic_model = BERTopic( embedding_model=embedding_model, 
+        topic_model = BERTopic( embedding_model=embedding_model_pipe, 
                                 #hdbscan_model=cluster_model,
                                 vectorizer_model=vectoriser_model,
                                 min_topic_size= min_docs_slider,
@@ -134,17 +169,18 @@ def extract_topics(in_files, in_file, min_docs_slider, in_colnames, max_topics_s
 
     # Do this if you have pre-assigned topics
     else:   
-        zero_shot_topics_list = read_file(candidate_topics.name)
-        zero_shot_topics_list_lower = [x.lower() for x in zero_shot_topics_list]
+        zero_shot_topics = read_file(candidate_topics.name)
+        #print(zero_shot_topics)
+        zero_shot_topics_lower = list(zero_shot_topics.iloc[:, 0].str.lower())
 
-        print(zero_shot_topics_list_lower)
+        print(zero_shot_topics_lower)
 
-        topic_model = BERTopic( embedding_model=embedding_model,
+        topic_model = BERTopic( embedding_model=embedding_model_pipe,
                                 #hdbscan_model=cluster_model,
                                 vectorizer_model=vectoriser_model,
                                 min_topic_size = min_docs_slider,
                                 nr_topics = max_topics_slider,
-                                zeroshot_topic_list = zero_shot_topics_list_lower,
+                                zeroshot_topic_list = zero_shot_topics_lower,
                                 zeroshot_min_similarity = 0.7,
                                 representation_model=representation_model,
                                 verbose = True)
@@ -152,44 +188,69 @@ def extract_topics(in_files, in_file, min_docs_slider, in_colnames, max_topics_s
         topics_text, probs = topic_model.fit_transform(docs, embeddings_out)
 
     if not topics_text:
-        return "No topics found, original file returned", data_file_name
+        return "No topics found, original file returned", data_file_name, None
         
     else: 
-        topics_text_out = topics_text
-        topics_scores_out = probs
-
-    topic_det_output_name = "topic_details_" + today_rev + ".csv"
+        print("Preparing topic model outputs.")
 
     topic_dets = topic_model.get_topic_info()
+    #print(topic_dets.columns)
 
-    topic_dets.to_csv(topic_det_output_name)
-    #print(topic_dets)
+    if topic_dets.shape[0] == 1:
+        topic_det_output_name = "topic_details_" + data_file_name_no_ext + "_" + today_rev + ".csv"
+        topic_dets.to_csv(topic_det_output_name)
 
-    doc_det_output_name = "doc_details_" + today_rev + ".csv"
-    doc_dets = topic_model.get_document_info(docs)[["Document",	"Topic", "Probability",	"Name", "Representative_document"]]
-    doc_dets.to_csv(doc_det_output_name)
-    #print(doc_dets)
-    
-    #print(topic_dets)    
-    #topics_text_out_str = ', '.join(list(topic_dets["KeyBERT"]))
+        return "No topics found, original file returned", [data_file_name, topic_det_output_name], None
 
-    topics_text_out_str = str(topic_dets["KeyBERT"])
-    #topics_scores_out_str = str(doc_dets["Probability"][0])
-    
-    output_text = "Topics: " + topics_text_out_str #+ "\n\nProbability scores: " + topics_scores_out_str
+    # Replace original labels with LLM labels
+    if "Mistral" in topic_model.get_topic_info().columns:
+        llm_labels = [label[0][0].split("\n")[0] for label in topic_model.get_topics(full=True)["Mistral"].values()]
+        topic_model.set_topic_labels(llm_labels)
+    else:
+        topic_model.set_topic_labels(list(topic_dets["Name"]))
 
     # Outputs
+    
+    topic_det_output_name = "topic_details_" + data_file_name_no_ext + "_" + today_rev + ".csv"
+    topic_dets.to_csv(topic_det_output_name)
+
+    doc_det_output_name = "doc_details_" + data_file_name_no_ext + "_" + today_rev + ".csv"
+    doc_dets = topic_model.get_document_info(docs)[["Document",	"Topic", "Name", "Representative_document"]] # "Probability",
+    doc_dets.to_csv(doc_det_output_name)
+
+    topics_text_out_str = str(topic_dets["Name"])
+    output_text = "Topics: " + topics_text_out_str
+   
     embedding_file_name = data_file_name_no_ext + '_' + 'embeddings.npz'
     np.savez_compressed(embedding_file_name, embeddings_out)
 
-    topic_model_save_name = data_file_name_no_ext + "_topics_" + today_rev + ".pkl"
-    topic_model.save(topic_model_save_name, serialization='pickle', save_embedding_model=False, save_ctfidf=False)
+    #if low_resource_mode == "No":
+    topic_model_save_name_folder = "output_model/" + data_file_name_no_ext + "_topics_" + today_rev# + ".safetensors"
+    topic_model_save_name_zip = topic_model_save_name_folder + ".zip"
+
+    # Clear folder before replacing files
+    delete_files_in_folder(topic_model_save_name_folder)
+
+    topic_model.save(topic_model_save_name_folder, serialization='safetensors', save_embedding_model=True, save_ctfidf=False)
+
+    # Zip file example
+    
+    zip_folder(topic_model_save_name_folder, topic_model_save_name_zip)
 
     # Visualise the topics:
     topics_vis = topic_model.visualize_documents(label_col, reduced_embeddings=reduced_embeddings, hide_annotations=True, hide_document_hover=False, custom_labels=True)
+
+    #return output_text, [doc_det_output_name, topic_det_output_name, embedding_file_name, topic_model_save_name_zip], topics_vis
+
+    #elif low_resource_mode == "Yes":
+    #    # Visualise the topics:
+    #    topics_vis = topic_model.visualize_documents(label_col, reduced_embeddings=reduced_embeddings, hide_annotations=True, hide_document_hover=False, custom_labels=True)
     
-    return output_text, [doc_det_output_name, topic_det_output_name, embedding_file_name, topic_model_save_name], topics_vis
+    #    return output_text, [doc_det_output_name, topic_det_output_name, embedding_file_name], topics_vis
+
+    return output_text, [doc_det_output_name, topic_det_output_name, embedding_file_name, topic_model_save_name_zip], topics_vis
 
+# , topic_model_save_name
 
 # ## Gradio app - extract topics
 
@@ -219,7 +280,7 @@ with block:
             candidate_topics = gr.File(label="Input topics from file (csv)")
             
         with gr.Row():
-            min_docs_slider = gr.Slider(minimum = 1, maximum = 1000, value = 15, step = 1, label = "Minimum number of documents needed to create topic")
+            min_docs_slider = gr.Slider(minimum = 2, maximum = 1000, value = 15, step = 1, label = "Minimum number of documents needed to create topic")
             max_topics_slider = gr.Slider(minimum = 2, maximum = 500, value = 3, step = 1, label = "Maximum number of topics")
 
         with gr.Row():
@@ -233,17 +294,19 @@ with block:
     
     with gr.Tab("Load and data processing options"):
         with gr.Accordion("Process data on load", open = True):
-            anonymise_drop = gr.Dropdown(value = "No", choices=["Yes", "No"], multiselect=False, label="Anonymise data on file load.")
-            return_intermediate_files = gr.Dropdown(label = "Return intermediate processing files from file preparation. Files can be loaded in to save processing time in future.", value="No", choices=["Yes", "No"])
-            embedding_super_compress = gr.Dropdown(label = "Round embeddings to three dp for smaller files with less accuracy.", value="No", choices=["Yes", "No"])
-            low_resource_mode_opt = gr.Dropdown(label = "Low resource mode (non-AI embeddings, no LLM-generated topic names).", value=low_resource_mode, choices=["Yes", "No"])
-                
+            with gr.Row():
+                anonymise_drop = gr.Dropdown(value = "No", choices=["Yes", "No"], multiselect=False, label="Anonymise data on file load. Names and other details are replaced with tags e.g. '<person>'.")
+                return_intermediate_files = gr.Dropdown(label = "Return intermediate processing files from file preparation. Files can be loaded in to save processing time in future.", value="No", choices=["Yes", "No"])
+                embedding_super_compress = gr.Dropdown(label = "Round embeddings to three dp for smaller files with less accuracy.", value="No", choices=["Yes", "No"])
+            with gr.Row():
+                low_resource_mode_opt = gr.Dropdown(label = "Low resource mode (non-AI embeddings, no LLM-generated topic names).", value="No", choices=["Yes", "No"])
+                create_llm_topic_labels = gr.Dropdown(label = "Create LLM-generated topic labels.", value="No", choices=["Yes", "No"])
 
     # Update column names dropdown when file uploaded
     in_files.upload(fn=put_columns_in_df, inputs=[in_files], outputs=[in_colnames, in_label, data_state])    
     in_colnames.change(dummy_function, in_colnames, None)
 
-    topics_btn.click(fn=extract_topics, inputs=[data_state, in_files, min_docs_slider, in_colnames, max_topics_slider, candidate_topics, in_label, anonymise_drop, return_intermediate_files, embedding_super_compress, low_resource_mode_opt], outputs=[output_single_text, output_file, plot], api_name="topics")
+    topics_btn.click(fn=extract_topics, inputs=[data_state, in_files, min_docs_slider, in_colnames, max_topics_slider, candidate_topics, in_label, anonymise_drop, return_intermediate_files, embedding_super_compress, low_resource_mode_opt, create_llm_topic_labels], outputs=[output_single_text, output_file, plot], api_name="topics")
 
 block.queue().launch(debug=True)#, server_name="0.0.0.0", ssl_verify=False, server_port=7860)
 

funcs/embeddings.py

@@ -11,7 +11,7 @@ if cuda.is_available():
 else: 
     torch_device =  "cpu"
 
-def make_or_load_embeddings(docs, file_list, data_file_name_no_ext, embedding_model, return_intermediate_files, embeddings_super_compress, low_resource_mode_opt):
+def make_or_load_embeddings(docs, file_list, data_file_name_no_ext, embedding_model, return_intermediate_files, embeddings_super_compress, low_resource_mode_opt, reduce_embeddings="Yes"):
 
     embeddings_file_names = [string.lower() for string in file_list if "embedding" in string.lower()]  
 
@@ -38,11 +38,19 @@ def make_or_load_embeddings(docs, file_list, data_file_name_no_ext, embedding_mo
             TruncatedSVD(100)
             )
 
-            embeddings_out = embedding_model.encode(sentences=docs, show_progress_bar = True, batch_size = 32)
+            # Fit the pipeline to the text data
+            embedding_model.fit(docs)
+
+            # Transform text data to embeddings
+            embeddings_out = embedding_model.transform(docs)
+
+            #embeddings_out = embedding_model.encode(sentences=docs, show_progress_bar = True, batch_size = 32)
 
         elif low_resource_mode_opt == "No":
             print("Creating dense embeddings based on transformers model")
 
+            #print("Embedding model is: ", embedding_model)
+
             embeddings_out = embedding_model.encode(sentences=docs, max_length=1024, show_progress_bar = True, batch_size = 32) # For Jina # # 
 
             #import torch
@@ -72,7 +80,8 @@ def make_or_load_embeddings(docs, file_list, data_file_name_no_ext, embedding_mo
                 np.savez_compressed(semantic_search_file_name, embeddings_out_round)
 
     # Pre-reduce embeddings for visualisation purposes
-    reduced_embeddings = UMAP(n_neighbors=15, n_components=2, min_dist=0.0, metric='cosine', random_state=42).fit_transform(embeddings_out)
-
+    if reduce_embeddings == "Yes":
+        reduced_embeddings = UMAP(n_neighbors=15, n_components=2, min_dist=0.0, metric='cosine', random_state=42).fit_transform(embeddings_out)
+        return embeddings_out, reduced_embeddings
 
-    return embeddings_out, reduced_embeddings
+    return embeddings_out, None

funcs/helper_functions.py

@@ -1,4 +1,5 @@
 import os
+import zipfile
 import re
 import pandas as pd
 import gradio as gr
@@ -87,3 +88,35 @@ def dummy_function(in_colnames):
     A dummy function that exists just so that dropdown updates work correctly.
     """
     return None
+
+# Zip the above to export file
+
+
+def zip_folder(folder_path, output_zip_file):
+    # Create a ZipFile object in write mode
+    with zipfile.ZipFile(output_zip_file, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        # Walk through the directory
+        for root, dirs, files in os.walk(folder_path):
+            for file in files:
+                # Create a complete file path
+                file_path = os.path.join(root, file)
+                # Add file to the zip file
+                # The arcname argument sets the archive name, i.e., the name within the zip file
+                zipf.write(file_path, arcname=os.path.relpath(file_path, folder_path))
+
+def delete_files_in_folder(folder_path):
+    # Check if the folder exists
+    if not os.path.exists(folder_path):
+        print(f"The folder {folder_path} does not exist.")
+        return
+
+    # Iterate over all files in the folder and remove each
+    for filename in os.listdir(folder_path):
+        file_path = os.path.join(folder_path, filename)
+        try:
+            if os.path.isfile(file_path) or os.path.islink(file_path):
+                os.unlink(file_path)
+            else:
+                print(f"Skipping {file_path} as it is a directory")
+        except Exception as e:
+            print(f"Failed to delete {file_path}. Reason: {e}")

funcs/prompts.py

@@ -37,7 +37,7 @@ ASSISTANT:Topic label:"""
 
 capybara_prompt = capybara_example_prompt + capybara_main_prompt
 
-print("Capybara prompt: ", capybara_prompt)
+#print("Capybara prompt: ", capybara_prompt)
 
 # System prompt describes information given to all conversations
 open_hermes_start="<|im_start|>"
@@ -72,7 +72,7 @@ Topic label:
 """
 open_hermes_prompt = open_hermes_system_prompt + open_hermes_example_prompt + open_hermes_main_prompt
 
-print("Open Hermes prompt: ", open_hermes_prompt)
+#print("Open Hermes prompt: ", open_hermes_prompt)
 
 stablelm_start = "<|user|>"
 stablelm_example_prompt = """<|user|>
@@ -103,4 +103,4 @@ Topic label:"""
 
 stablelm_prompt = stablelm_example_prompt + stablelm_main_prompt
 
-print("StableLM prompt: ", stablelm_prompt)
+#print("StableLM prompt: ", stablelm_prompt)

funcs/representation_model.py

@@ -9,8 +9,6 @@ import torch.cuda
 from bertopic.representation import KeyBERTInspired, MaximalMarginalRelevance, TextGeneration
 from funcs.prompts import capybara_prompt, capybara_start, open_hermes_prompt, open_hermes_start, stablelm_prompt, stablelm_start
 
-#from huggingface_hub import hf_hub_download
-#hf_hub_download(repo_id='second-state/stablelm-2-zephyr-1.6b-GGUF', filename='stablelm-2-zephyr-1_6b-Q5_K_M.gguf')
 
 hf_model_name =  'TheBloke/phi-2-orange-GGUF' #'NousResearch/Nous-Capybara-7B-V1.9-GGUF' # 'second-state/stablelm-2-zephyr-1.6b-GGUF'
 hf_model_file =   'phi-2-orange.Q5_K_M.gguf' #'Capybara-7B-V1.9-Q5_K_M.gguf' # 'stablelm-2-zephyr-1_6b-Q5_K_M.gguf'
@@ -18,9 +16,9 @@ chosen_prompt = open_hermes_prompt # stablelm_prompt
 chosen_start_tag =  open_hermes_start # stablelm_start
 
 # Find model file
-def find_model_file(hf_model_name, hf_model_file):
-    hf_loc = os.environ["HF_HOME"]
-    hf_sub_loc = os.environ["HF_HOME"] + "/hub/"
+def find_model_file(hf_model_name, hf_model_file, search_folder):
+    hf_loc = search_folder #os.environ["HF_HOME"]
+    hf_sub_loc = search_folder + "/hub/" #os.environ["HF_HOME"] 
 
     hf_model_name_path = hf_sub_loc + 'models--' + hf_model_name.replace("/","--")
 
@@ -36,16 +34,19 @@ def find_model_file(hf_model_name, hf_model_file):
     folder_path = hf_model_name_path  # Replace with your folder path
     file_to_find = hf_model_file         # Replace with the file name you're looking for
 
-    found_file = find_file(folder_path, file_to_find)
+    found_file = find_file(folder_path, file_to_find) # os.environ["HF_HOME"]
     if found_file:
         print(f"File found: {found_file}")
         return found_file
     else:
         error = "File not found."
-        print(error)
-        return error
+        print(error, " Downloading model from hub")
+        from huggingface_hub import hf_hub_download
+        hf_hub_download(repo_id=hf_model_name, filename='phi-2-orange.Q5_K_M.gguf')
+        found_file = find_file(folder_path, file_to_find)
+        return found_file
 
-found_file = find_model_file(hf_model_name, hf_model_file)
+found_file = find_model_file(hf_model_name, hf_model_file, os.environ["HF_HOME"])#".")
 
 # Currently set n_gpu_layers to 0 even with cuda due to persistent bugs in implementation with cuda
 if torch.cuda.is_available():
@@ -57,7 +58,7 @@ else:
     low_resource_mode = "Yes"
     n_gpu_layers = 0
 
-#low_resource_mode = "Yes"
+low_resource_mode = "No" # Override for testing
 
 #print("Running on device:", torch_device)
 n_threads = torch.get_num_threads()
@@ -140,32 +141,32 @@ gen_config = LLamacppGenerateConfig(
 # KeyBERT
 keybert = KeyBERTInspired()
 
-if low_resource_mode == "No":
-    # Use llama.cpp to load in model
-    llm = Llama(model_path=found_file, stop=chosen_start_tag, n_gpu_layers=n_gpu_layers, n_ctx=n_ctx) #**gpu_config.model_dump())# 
-    #print(llm.n_gpu_layers)
-    llm_model = LlamaCPP(llm, prompt=chosen_prompt)#, **gen_config.model_dump())
-
-    # All representation models
-    representation_model = {
-    "KeyBERT": keybert,
-    "Mistral": llm_model
-    }
+def create_representation_model(create_llm_topic_labels, gpu_config, found_file, chosen_start_tag):
 
-elif low_resource_mode == "Yes":
-    representation_model = {"KeyBERT": keybert}
+    if create_llm_topic_labels == "Yes":
+        # Use llama.cpp to load in model
+        llm = Llama(model_path=found_file, stop=chosen_start_tag, n_gpu_layers=gpu_config.n_gpu_layers, n_ctx=gpu_config.n_ctx) #**gpu_config.model_dump())# 
+        #print(llm.n_gpu_layers)
+        llm_model = LlamaCPP(llm, prompt=chosen_prompt)#, **gen_config.model_dump())
 
-# Deprecated example using CTransformers. This package is not really used anymore
-#model = AutoModelForCausalLM.from_pretrained('NousResearch/Nous-Capybara-7B-V1.9-GGUF', model_type='mistral', model_file='Capybara-7B-V1.9-Q5_K_M.gguf', hf=True, **vars(gpu_config))
-#tokenizer = AutoTokenizer.from_pretrained("NousResearch/Nous-Capybara-7B-V1.9")
-#generator = pipeline(task="text-generation", model=model, tokenizer=tokenizer)
+        # All representation models
+        representation_model = {
+        "KeyBERT": keybert,
+        "Mistral": llm_model
+        }
 
-# Text generation with Llama 2
-#mistral_capybara = TextGeneration(generator, prompt=capybara_prompt)
-#mistral_hermes = TextGeneration(generator, prompt=open_hermes_prompt)
+    elif create_llm_topic_labels == "No":
+        representation_model = {"KeyBERT": keybert}
 
+    # Deprecated example using CTransformers. This package is not really used anymore
+    #model = AutoModelForCausalLM.from_pretrained('NousResearch/Nous-Capybara-7B-V1.9-GGUF', model_type='mistral', model_file='Capybara-7B-V1.9-Q5_K_M.gguf', hf=True, **vars(gpu_config))
+    #tokenizer = AutoTokenizer.from_pretrained("NousResearch/Nous-Capybara-7B-V1.9")
+    #generator = pipeline(task="text-generation", model=model, tokenizer=tokenizer)
 
+    # Text generation with Llama 2
+    #mistral_capybara = TextGeneration(generator, prompt=capybara_prompt)
+    #mistral_hermes = TextGeneration(generator, prompt=open_hermes_prompt)
+        
+    return representation_model
 
-# MMR (is rubbish, don't use)
-#mmr = MaximalMarginalRelevance(diversity=0.3)