Stop generation button. Better model load. Trying one source that's longer.

app.py

@@ -13,7 +13,6 @@ from langchain.vectorstores import FAISS
 import gradio as gr
 
 from transformers import AutoTokenizer
-from dataclasses import asdict, dataclass
 
 # Alternative model sources
 from ctransformers import AutoModelForCausalLM
@@ -83,7 +82,7 @@ def load_model(model_type, gpu_layers, gpu_config=None, cpu_config=None, torch_d
     if model_type == "Orca Mini":
 
         gpu_config.update_gpu(gpu_layers)
-        cpu_config.update_gpu(0)
+        cpu_config.update_gpu(gpu_layers)
 
         print("Loading with", cpu_config.gpu_layers, "model layers sent to GPU.")
 
@@ -92,8 +91,13 @@ def load_model(model_type, gpu_layers, gpu_config=None, cpu_config=None, torch_d
 
         try:
             model = AutoModelForCausalLM.from_pretrained('juanjgit/orca_mini_3B-GGUF', model_type='llama', model_file='orca-mini-3b.q4_0.gguf', **vars(gpu_config)) # **asdict(CtransRunConfig_cpu())
+            #model = AutoModelForCausalLM.from_pretrained('Aryanne/Sheared-LLaMA-1.3B-gguf', model_type='llama', model_file='q8_0-sheared-llama-1.3b.gguf', **vars(gpu_config)) # **asdict(CtransRunConfig_cpu())
+            #model = AutoModelForCausalLM.from_pretrained('TheBloke/TinyLlama-1.1B-1T-OpenOrca-GGUF', model_type='llama', model_file='tinyllama-1.1b-1t-openorca.Q8_0.gguf', **vars(gpu_config)) # **asdict(CtransRunConfig_cpu())
         except:
             model = AutoModelForCausalLM.from_pretrained('juanjgit/orca_mini_3B-GGUF', model_type='llama', model_file='orca-mini-3b.q4_0.gguf', **vars(cpu_config)) #**asdict(CtransRunConfig_gpu())
+            #model = AutoModelForCausalLM.from_pretrained('Aryanne/Sheared-LLaMA-1.3B-gguf', model_type='llama', model_file='q8_0-sheared-llama-1.3b.gguf', **vars(gpu_config)) # **asdict(CtransRunConfig_cpu())
+            #model = AutoModelForCausalLM.from_pretrained('TheBloke/TinyLlama-1.1B-1T-OpenOrca-GGUF', model_type='llama', model_file='tinyllama-1.1b-1t-openorca.Q8_0.gguf', **vars(cpu_config)) # **asdict(CtransRunConfig_cpu())
+
 
         tokenizer = []
 
@@ -126,8 +130,10 @@ def load_model(model_type, gpu_layers, gpu_config=None, cpu_config=None, torch_d
     chatf.tokenizer = tokenizer
     chatf.model_type = model_type
 
-    print("Finished loading model: ", model_type)
-    return model_type
+    load_confirmation = "Finished loading model: " + model_type
+
+    print(load_confirmation)
+    return model_type, load_confirmation
 
 # Both models are loaded on app initialisation so that users don't have to wait for the models to be downloaded
 model_type = "Orca Mini"
@@ -173,7 +179,7 @@ with block:
 
     gr.Markdown("<h1><center>Lightweight PDF / web page QA bot</center></h1>")        
     
-    gr.Markdown("Chat with PDF or web page documents. The default is a small model (Flan Alpaca), that can only answer specific questions that are answered in the text. It cannot give overall impressions of, or summarise the document. The alternative (Orca Mini), can reason a little better, but is much slower (See Advanced tab).\n\nBy default the Lambeth Borough Plan '[Lambeth 2030 : Our Future, Our Lambeth](https://www.lambeth.gov.uk/better-fairer-lambeth/projects/lambeth-2030-our-future-our-lambeth)' is loaded. If you want to talk about another document or web page, please select from the second tab. If switching topic, please click the 'Clear chat' button.\n\nCaution: This is a public app. Likes and dislike responses will be saved to disk to improve the model. Please ensure that the document you upload is not sensitive is any way as other users may see it! Also, please note that LLM chatbots may give incomplete or incorrect information, so please use with care.")
+    gr.Markdown("Chat with PDF or web page documents. The default is a small model (Flan Alpaca), that can only answer specific questions that are answered in the text. It cannot give overall impressions of, or summarise the document. The alternative (Orca Mini), can reason a little better, but is much slower (See Advanced tab).\n\nBy default the Lambeth Borough Plan '[Lambeth 2030 : Our Future, Our Lambeth](https://www.lambeth.gov.uk/better-fairer-lambeth/projects/lambeth-2030-our-future-our-lambeth)' is loaded. If you want to talk about another document or web page, please select from the second tab. If switching topic, please click the 'Clear chat' button.\n\nCaution: This is a public app. Please ensure that the document you upload is not sensitive is any way as other users may see it! Also, please note that LLM chatbots may give incomplete or incorrect information, so please use with care.")
 
     current_source = gr.Textbox(label="Current data source that is loaded into the app", value="Lambeth_2030-Our_Future_Our_Lambeth.pdf")
 
@@ -181,8 +187,8 @@ with block:
 
         with gr.Row():
             chat_height = 500
-            chatbot = gr.Chatbot(height=chat_height, avatar_images=('user.jfif', 'bot.jpg'),bubble_full_width = False)
-            sources = gr.HTML(value = "Source paragraphs where I looked for answers will appear here", height=chat_height)
+            chatbot = gr.Chatbot(height=chat_height, avatar_images=('user.jfif', 'bot.jpg'),bubble_full_width = False, scale = 1)
+            sources = gr.HTML(value = "Source paragraphs where I looked for answers will appear here", height=chat_height, scale = 2)
 
         with gr.Row():
             message = gr.Textbox(
@@ -191,7 +197,8 @@ with block:
             )     
         with gr.Row():
             submit = gr.Button(value="Send message", variant="secondary", scale = 1)
-            clear = gr.Button(value="Clear chat", variant="secondary", scale=0)  
+            clear = gr.Button(value="Clear chat", variant="secondary", scale=0) 
+            stop = gr.Button(value="Stop generating", variant="secondary", scale=0)
 
         examples_set = gr.Radio(label="Examples for the Lambeth Borough Plan",
             #value = "What were the five pillars of the previous borough plan?",
@@ -220,7 +227,10 @@ with block:
 
     with gr.Tab("Advanced features"):
         model_choice = gr.Radio(label="Choose a chat model", value="Flan Alpaca", choices = ["Flan Alpaca", "Orca Mini"])
-        gpu_layer_choice = gr.Slider(label="Choose number of model layers to send to GPU (please don't change if you don't know what you're doing).", value=0, minimum=0, maximum=6, step = 1, scale = 0, visible=False)
+        with gr.Row():
+            gpu_layer_choice = gr.Slider(label="Choose number of model layers to send to GPU (WARNING: please don't modify unless you have a GPU).", value=0, minimum=0, maximum=6, step = 1, visible=True)
+            change_model_button = gr.Button(value="Load model", scale=0)
+        load_text = gr.Text(label="Load status")
 
     gr.HTML(
         "<center>This app is based on the models Flan Alpaca and Orca Mini. It powered by Gradio, Transformers, Ctransformers, and Langchain.</a></center>"
@@ -228,7 +238,11 @@ with block:
 
     examples_set.change(fn=chatf.update_message, inputs=[examples_set], outputs=[message])
 
-    model_choice.change(fn=load_model, inputs=[model_choice, gpu_layer_choice], outputs = [model_type_state])
+    change_model_button.click(fn=chatf.turn_off_interactivity, inputs=[message, chatbot], outputs=[message, chatbot], queue=False).\
+    then(fn=load_model, inputs=[model_choice, gpu_layer_choice], outputs = [model_type_state, load_text]).\
+    then(lambda: chatf.restore_interactivity(), None, [message], queue=False).\
+    then(chatf.clear_chat, inputs=[chat_history_state, sources, message, current_topic], outputs=[chat_history_state, sources, message, current_topic]).\
+    then(lambda: None, None, chatbot, queue=False)
 
     # Load in a pdf
     load_pdf_click = load_pdf.click(ing.parse_file, inputs=[in_pdf], outputs=[ingest_text, current_source]).\
@@ -259,6 +273,9 @@ with block:
                 then(chatf.add_inputs_answer_to_history,[message, chatbot, current_topic], [chat_history_state, current_topic]).\
                 then(lambda: chatf.restore_interactivity(), None, [message], queue=False)
     
+    # Stop box
+    stop.click(fn=None, inputs=None, outputs=None, cancels=[response_click, response_enter])
+
     # Clear box
     clear.click(chatf.clear_chat, inputs=[chat_history_state, sources, message, current_topic], outputs=[chat_history_state, sources, message, current_topic])
     clear.click(lambda: None, None, chatbot, queue=False)

chatfuncs/chatfuncs.py

@@ -69,7 +69,7 @@ kw_model = pipeline("feature-extraction", model="sentence-transformers/all-MiniL
 
 if torch.cuda.is_available():
     torch_device = "cuda"
-    gpu_layers = 6
+    gpu_layers = 0
 else: 
     torch_device =  "cpu"
     gpu_layers = 0
@@ -82,25 +82,38 @@ print("CPU threads:", threads)
 temperature: float = 0.1
 top_k: int = 3
 top_p: float = 1
-repetition_penalty: float = 1.05
+repetition_penalty: float = 1.3
 flan_alpaca_repetition_penalty: float = 1.3
+tinyllama_repetition_penalty: float = 1.5
 last_n_tokens: int = 64
-max_new_tokens: int = 125
+max_new_tokens: int = 512
 seed: int = 42
 reset: bool = False
 stream: bool = True
 threads: int = threads
-batch_size:int = 1024
-context_length:int = 4096
+batch_size:int = 256
+context_length:int = 2048
 sample = True
 
 
 class CtransInitConfig_gpu:
-    def __init__(self, temperature=0.1, top_k=3, top_p=1, repetition_penalty=1.05, last_n_tokens=64, max_new_tokens=125, seed=42, reset=False, stream=True, threads=None, batch_size=1024, context_length=4096, gpu_layers=None):
+    def __init__(self, temperature=temperature,
+                 top_k=top_k,
+                 top_p=top_p,
+                 repetition_penalty=repetition_penalty,
+                 last_n_tokens=last_n_tokens,
+                 max_new_tokens=max_new_tokens,
+                 seed=seed,
+                 reset=reset,
+                 stream=stream,
+                 threads=threads,
+                 batch_size=batch_size,
+                 context_length=context_length,
+                 gpu_layers=gpu_layers):
         self.temperature = temperature
         self.top_k = top_k
         self.top_p = top_p
-        self.repetition_penalty = repetition_penalty
+        self.repetition_penalty = repetition_penalty# repetition_penalty
         self.last_n_tokens = last_n_tokens
         self.max_new_tokens = max_new_tokens
         self.seed = seed
@@ -124,17 +137,38 @@ gpu_config = CtransInitConfig_gpu()
 cpu_config = CtransInitConfig_cpu()
 
 
-@dataclass
+#@dataclass
+#class CtransGenGenerationConfig:
+#    top_k: int = top_k
+#    top_p: float = top_p
+#    temperature: float = temperature
+#    repetition_penalty: float = tinyllama_repetition_penalty
+#    last_n_tokens: int = last_n_tokens
+#    seed: int = seed
+#    batch_size:int = batch_size
+#     threads: int = threads
+#     reset: bool = True
+
 class CtransGenGenerationConfig:
-    top_k: int = top_k
-    top_p: float = top_p
-    temperature: float = temperature
-    repetition_penalty: float = repetition_penalty
-    last_n_tokens: int = last_n_tokens
-    seed: int = seed
-    batch_size:int = batch_size
-    threads: int = threads
-    reset: bool = True
+    def __init__(self, temperature=temperature,
+                 top_k=top_k,
+                 top_p=top_p,
+                 repetition_penalty=repetition_penalty,
+                 last_n_tokens=last_n_tokens,
+                 seed=seed,
+                 threads=threads,
+                 batch_size=batch_size,
+                 reset=True
+                 ):
+        self.temperature = temperature
+        self.top_k = top_k
+        self.top_p = top_p
+        self.repetition_penalty = repetition_penalty# repetition_penalty
+        self.last_n_tokens = last_n_tokens
+        self.seed = seed
+        self.threads = threads
+        self.batch_size = batch_size
+        self.reset = reset
 
 # Vectorstore funcs
 
@@ -199,6 +233,12 @@ def base_prompt_templates(model_type = "Flan Alpaca"):
 
     Response:"""
 
+    instruction_prompt_template_sheared_llama = """Answer the QUESTION using information from the following CONTENT.
+    CONTENT: {summaries}
+    QUESTION: {question}
+
+    Answer:"""
+
     instruction_prompt_template_orca = """
     ### System:
     You are an AI assistant that follows instruction extremely well. Help as much as you can.
@@ -215,7 +255,15 @@ def base_prompt_templates(model_type = "Flan Alpaca"):
     Answer the QUESTION using information from the following CONTENT. Respond with short answers that directly answer the question.
     CONTENT: {summaries}
     QUESTION: {question}\n
-    <|im_end|>"""
+    Answer:<|im_end|>"""
+
+    instruction_prompt_tinyllama_orca = """<|im_start|>system\n
+    You are an AI assistant that follows instruction extremely well. Help as much as you can.
+    <|im_start|>user\n
+    Answer the QUESTION using information from the following CONTENT. Only quote text that directly answers the question and nothing more. If you can't find an answer to the question, respond with "Sorry, I can't find an answer to that question.".
+    CONTENT: {summaries}
+    QUESTION: {question}\n
+    Answer:<|im_end|>"""
 
     if model_type == "Flan Alpaca":
         INSTRUCTION_PROMPT=PromptTemplate(template=instruction_prompt_template_alpaca, input_variables=['question', 'summaries'])
@@ -233,12 +281,12 @@ def generate_expanded_prompt(inputs: Dict[str, str], instruction_prompt, content
         new_question_kworded = adapt_q_from_chat_history(question, chat_history, extracted_memory) # new_question_keywords, 
         
        
-        docs_keep_as_doc, doc_df, docs_keep_out = hybrid_retrieval(new_question_kworded, vectorstore, embeddings, k_val = 5, out_passages = 2,
+        docs_keep_as_doc, doc_df, docs_keep_out = hybrid_retrieval(new_question_kworded, vectorstore, embeddings, k_val = 5, out_passages = 1,
                                                                           vec_score_cut_off = 1, vec_weight = 1, bm25_weight = 1, svm_weight = 1)#,
                                                                           #vectorstore=globals()["vectorstore"], embeddings=globals()["embeddings"])
         
         # Expand the found passages to the neighbouring context
-        docs_keep_as_doc, doc_df = get_expanded_passages(vectorstore, docs_keep_out, width=1)
+        docs_keep_as_doc, doc_df = get_expanded_passages(vectorstore, docs_keep_out, width=3)
 
         if docs_keep_as_doc == []:
             {"answer": "I'm sorry, I couldn't find a relevant answer to this question.", "sources":"I'm sorry, I couldn't find a relevant source for this question."}
@@ -301,7 +349,7 @@ def produce_streaming_answer_chatbot(history, full_prompt, model_type):
             streamer=streamer,
             max_new_tokens=max_new_tokens,
             do_sample=sample,
-            repetition_penalty=flan_alpaca_repetition_penalty,
+            repetition_penalty=repetition_penalty,
             top_p=top_p,
             temperature=temperature,
             top_k=top_k
@@ -332,13 +380,15 @@ def produce_streaming_answer_chatbot(history, full_prompt, model_type):
     elif model_type == "Orca Mini":
         tokens = model.tokenize(full_prompt)
 
+        gen_config = CtransGenGenerationConfig()
+
         # Pull the generated text from the streamer, and update the model output.
         start = time.time()
         NUM_TOKENS=0
         print('-'*4+'Start Generation'+'-'*4)
 
         history[-1][1] = ""
-        for new_text in model.generate(tokens, **asdict(CtransGenGenerationConfig())): #CtransGen_generate(prompt=full_prompt)#, config=CtransGenGenerationConfig()): # #top_k=top_k, temperature=temperature, repetition_penalty=repetition_penalty,
+        for new_text in model.generate(tokens, **vars(gen_config)): #CtransGen_generate(prompt=full_prompt)#, config=CtransGenGenerationConfig()): # #top_k=top_k, temperature=temperature, repetition_penalty=repetition_penalty,
             if new_text == None: new_text =  ""
             history[-1][1] += model.detokenize(new_text) #new_text
             NUM_TOKENS+=1