Upload 98 files

.gitattributes

@@ -0,0 +1,2 @@
+extensions/superboogav2/nltk_data/corpora/wordnet.zip filter=lfs diff=lfs merge=lfs -text
+extensions/superboogav2/nltk_data/taggers/averaged_perceptron_tagger/averaged_perceptron_tagger.pickle filter=lfs diff=lfs merge=lfs -text

extensions/Training_PRO/README.md

@@ -0,0 +1,92 @@
+# Training_PRO
+
+This is an expanded and reworked Training tab
+Maintained by FP
+
+[![ko-fi](https://ko-fi.com/img/githubbutton_sm.svg)](https://ko-fi.com/Q5Q5MOB4M)
+
+Repo home:
+
+https://github.com/FartyPants/Training_PRO
+
+In general the repo above is ahead of the extension included in text WebUi.
+
+## News
+
+- NEFtune: add noise to help with generalization
+- Loss Graph in interface.
+- Supports Mistral training
+- some roundabout around pytorch and transformers version desync
+
+![image](https://github.com/FartyPants/Training_PRO/assets/23346289/e389ec69-d7ad-4922-9ad9-865625997479)
+
+## Features/Changes
+
+- Chunking: precise raw text slicer (PRTS) uses sentence slicing and making sure things are clean on all ends
+- overlap chunking - this special overlapping will make additional overlap block based on logical rules (aka no overlap block on hard cut)
+- custom scheduler (follow the code to make your own) In LR Scheduler select FP_low_epoch_annealing - this scheduler will keep the LR constant for first epoch then use cosine for the rest - this part would be best to spawn into a new py file
+- saves graph png file at the end with learning rate and loss per epoch
+- adding EOS to each block or to hard cut only
+- automatically lowers gradient accumulation if you go overboard and set gradient accumulation that will be higher than actual data - transformers would then throw error (or they used to, not sure if still true) but in any way, it will fix bad data
+- turn BOS on and OFF
+- target selector
+- DEMENTOR LEARNING (experimental) Deep Memorization Enforcement Through Overlapping and Repetition. This is an experiment for long-text learning using low epochs (basically use 1 epoch with constant LR or 2 epochs with FP_low_epoch_annealing LR scheduler)
+- Getting rid of micro batch size/batch size confusion. Now there is True Batch Size and Gradient accumulation slider, consisten with all the other training out there
+- Ability to save Checkpoint during training with a button
+- Ability to change Stop Loss during training
+- different modes of checkpoint auto saving
+- Function to Check Dataset and suggest parameters such as warmup and checkpoint save frequency before training
+- Graph Training Loss in interface
+- more custom schedulers
+  
+### Notes:
+
+This uses it's own chunking code for raw text based on sentence splitting. This will avoid weird cuts in the chunks and each chunk should now start with sentence and end on some sentence. It works hand in hand with Hard Cut. A propper use is to structure your text into logical blocks (ideas) separated by three \n then use three \n in hard cut. This way each chunk will contain only one flow of ideas and not derail in the thoughts. And Overlapping code will create overlapped blocks on sentence basis too, but not cross hard cut, thus not cross different ideas either. Does it make any sense? No? Hmmmm...
+
+### Custom schedulers
+
+A bunch of custom (combination) schedulers are added to the LR schedule. These are based on my own experiments
+
+**FP_low_epoch_annealing**
+
+Uses constant LR (with warmup) for 1 epoch only. The rest of the epoch(s) is cosine annealing. So 10 epochs - 1 will be constant 9 will be nose dive down. However a typical usage would be 2 epochs (hence low epoch in name). 1st is constant, the second is annealing. Simple. I use it 90% of time.
+
+**FP_half_time_annealing**
+
+Like the low epoch, but now the total number of steps is divided by 2. First half is constant, second half is annealing. So 10 epochs - 5 will be constant, 5 will be cosine nose down.
+
+**FP_raise_fall_creative**
+
+This is a sine raise till half of the total steps then cosine fall the rest. (Or you may think of the curve as sine in its entirety. The most learning is done in the hump, in the middle. The warmup entry has no effect, since sine is automatically warm up.
+The idea is to start very mildly as not to overfit with the first blocks of dataset. It seems to broaden the scope of the model making it less strict for tight dataset. 
+
+### Targets
+
+Normal LORA is q, v and that's what you should use. You can use (q k v o) or (q k v) and it will give you a lot more trainable parameters. The benefit is that you can keep rank lower and still attain the same coherency as q v with high rank. Guanaco has been trained with QLORA and q k v o for example and they swear by it.
+
+### DEMENTOR LEARNING (experimental) Deep Memorization Enforcement Through Overlapping and Repetition
+
+This is and experimental chunking to train long-form text in low number of epochs (basically 1) with sliding repetition. The depth of learning directly depends on the cutoff_length. Increasing cutoff length will also increase number of blocks created from long-form text (which is contrary to normal training). It is based on my own wild experiments. 
+
+### Getting rid of batch size and micro batch size
+
+Keeping consistency with everyone else. 
+
+Listen, There is only ONE batch size - the True batch size (called previously micro-batch size in WebUI) - this is how many blocks are processed at once (during a single step). It eats GPU, but it really helps with the quality training (in fact the ideal batch size would be the same as number of blocks - which is unrealistic) - so the idea is to cram as much True Batch Size before your GPU blows with OOM. On 24GB this is about 10 for 13b (loaded with 4-bit)
+
+So no micro batch size - it is now called True Batch Size, because that's what it is.
+
+The other thing is Gradient Accumulation - this is an emulation of the above Batch size - a virtual batch size, if you will. If your GPU can't handle real batch size then you may fake it using Gradient Accumulation. This will accumulate the gradients over so many steps defined here and then update the weights at the end without increase in GPU.
+Gradient accumulation is like a virtual Batch size multiplier without the GPU penalty.
+
+If your batch size is 4 and your gradient accumulation is 2 then it sort of behaves as if we have batch size 8. *Sort of* because Batch size of 4 and GA of 2 is NOT the same as batch size of 2 and GA of 4. (It produces different weights - hence it's not an equivalent). The idea is that if you don't have GPU - using GA to extend batch size is the next best thing (good enough) since you have no other choice.
+
+If all you can afford is 1 batch size, then increasing GA will likely make the learning better in some range of GA (it's not always more is better).
+
+However - GA is not some golden goose. As said, it isn't the same as batch size. In fact GA may worsen your learning as well.
+
+I would suggest a series of experiment where you would put batch size as high as possible without OOM, set GA 1, then repeat training while increasing the GA (2, 4...), and see how the model changes. It's likely that it would follow some sort of curve where GA will seem to help before it will make it worse. Some people believe that if you can squeeze 6 BATCH Size, then you should not bother with GA at all... YMMW
+
+High Batch Size vs High GA would also likely produce different results in terms of learning  words vs style. How? Hmmmm... good question.
+
+One optical "benefit" of GA is that the loss will fluctuate less (because of all the gradient accumulation, which works as a form of noise smoothing as well).

extensions/Training_PRO/custom_scheduler.py

@@ -0,0 +1,433 @@
+from functools import partial
+import torch
+import transformers
+import math
+from torch.optim.lr_scheduler import LambdaLR
+
+from peft import (
+    PeftModel,
+)
+
+RED = "\033[91m"
+YELLOW = "\033[93m"
+GREEN = "\033[92m"
+RESET = "\033[0m"
+
+last_print_label = ''
+
+custom_scheduler_params = {'trigger_loss': 0.0, 'ramp_down_ratio':1.0, 'current_loss': 0.0,'dynamic_scheduler_stop': False, 'calc_ramp_down_at_step': 0, 'calc_num_training_steps': 0}
+
+
+def custom_scheduler_global_update(current_loss: float):
+    custom_scheduler_params.update({'current_loss': current_loss})
+  
+def custom_scheduler_global_setup(trigger_loss: float, ramp_down_ratio: float):
+    custom_scheduler_params.update({'trigger_loss': trigger_loss})
+    custom_scheduler_params.update({'ramp_down_ratio': ramp_down_ratio})
+
+    # calculates the total num steps after trigger
+    custom_scheduler_params.update({'calc_num_training_steps': 0})
+    #calculates steps when the ramp_down trigger occured
+    custom_scheduler_params.update({'calc_ramp_down_at_step': 0})
+    # triggers scheduler stopping after it reached calc_num_training_steps
+    custom_scheduler_params.update({'dynamic_scheduler_stop': False})
+
+
+# hold constant to the half of epochs then cosine down to 0
+def _get_fp_half_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_firstepoch_steps: int):
+    
+    global last_print_label
+    print_label = ''
+
+    half_steps = num_training_steps//2
+    
+    num_warmup_steps = min(num_warmup_steps,half_steps)
+
+    if current_step < num_warmup_steps:
+        print_label = 'Scheduler: Warmup'
+    elif current_step < half_steps:
+        print_label = 'Scheduler: Hold'
+    else:
+        print_label = 'Scheduler: Annealing'
+    
+    if print_label != last_print_label:
+        print(print_label)
+    
+    last_print_label = print_label
+
+    if current_step < num_warmup_steps:
+        return float(current_step) / float(max(1, num_warmup_steps))
+    
+    if current_step < half_steps:
+        return 1.0 
+    
+    progress = float(current_step - half_steps) / float(max(1, num_training_steps - half_steps))
+    num_cycles = 0.5
+    return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))    
+ 
+
+# raise up in cosine, then fall back in cosine
+def _get_fp_cosine_raise_and_fall_lr_lambda(current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_firstepoch_steps: int):
+    
+    global last_print_label
+    print_label = ''
+
+    half_steps = num_training_steps//2
+    
+    #num_warmup_steps = min(num_warmup_steps,half_steps)
+
+    if current_step < half_steps:
+        print_label = 'Scheduler: Raise'
+    else:
+        print_label = 'Scheduler: Fall'
+    
+    if print_label != last_print_label:
+        print(print_label)
+    
+    last_print_label = print_label
+
+    
+    # linear
+    #    return float(current_step) / float(max(1, num_warmup_steps))
+    
+    progress = float(current_step - half_steps) / float(max(1, num_training_steps - half_steps))
+    num_cycles = 0.5
+    return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))    
+ 
+# constant to the first epochs then cosine down to 0 over the rest epochs
+def _get_fp_cosine_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_firstepoch_steps: int):
+    
+    global last_print_label
+    print_label = ''
+    
+    num_warmup_steps = min(num_warmup_steps,num_firstepoch_steps)
+
+    if current_step < num_warmup_steps:
+        print_label = 'Scheduler: Warmup'
+    elif current_step < num_firstepoch_steps:
+        print_label = 'Scheduler: Hold'
+    else:
+        print_label = 'Scheduler: Annealing'
+    
+    if print_label != last_print_label:
+        print(print_label)
+    
+    last_print_label = print_label
+
+    if current_step < num_warmup_steps:
+        return float(current_step) / float(max(1, num_warmup_steps))
+    
+    if current_step < num_firstepoch_steps:
+        return 1.0 
+    
+    progress = float(current_step - num_firstepoch_steps) / float(max(1, num_training_steps - num_firstepoch_steps))
+    num_cycles = 0.5
+    return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))    
+    
+# halve lr each epoch   
+
+def _get_fp_cdrop_rate_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_firstepoch_steps: int):
+    
+    global last_print_label
+    print_label = ''
+    
+    num_warmup_steps = min(num_warmup_steps, num_firstepoch_steps)
+
+    current_epoch = (current_step // num_firstepoch_steps) + 1
+    
+    
+    if current_step < num_warmup_steps:
+        print_label = 'Scheduler: Warmup'
+    elif current_step < num_firstepoch_steps:
+        print_label = 'Scheduler: Hold'
+    else:
+        print_label = 'Scheduler: Drop Rate'
+    
+    if print_label != last_print_label:
+        print(print_label)
+    
+    last_print_label = print_label
+
+    if current_step < num_warmup_steps:
+        return float(current_step) / float(max(1, num_warmup_steps))
+    
+    if current_step < num_firstepoch_steps:
+        return 1.0 
+
+    # Compute the learning rate for the annealing phase
+    
+    learning_rate = 1.0 / float(2 ** (current_epoch - 1))
+   
+    return learning_rate
+
+# epoch decay: 1/(1 + decay * epoch)
+
+def custom_cosine_scheduler_with_warmup(optimizer, num_warmup_steps, num_training_steps, num_firstepoch_steps, last_epoch=-1):
+    """
+    Args:
+        optimizer ([`~torch.optim.Optimizer`]):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (`int`):
+            The number of steps for the warmup phase.
+        num_training_steps (`int`):
+            The total number of training steps.
+        last_epoch (`int`, *optional*, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+    """
+    
+    lr_lambda = partial(
+        _get_fp_cosine_schedule_with_warmup_lr_lambda,
+        num_warmup_steps=num_warmup_steps,
+        num_training_steps=num_training_steps,
+        num_firstepoch_steps = num_firstepoch_steps,
+    )
+    return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+def custom_half_scheduler_with_warmup(optimizer, num_warmup_steps, num_training_steps, num_firstepoch_steps, last_epoch=-1):
+    """
+    Args:
+        optimizer ([`~torch.optim.Optimizer`]):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (`int`):
+            The number of steps for the warmup phase.
+        num_training_steps (`int`):
+            The total number of training steps.
+        last_epoch (`int`, *optional*, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+    """
+    
+    lr_lambda = partial(
+        _get_fp_half_schedule_with_warmup_lr_lambda,
+        num_warmup_steps=num_warmup_steps,
+        num_training_steps=num_training_steps,
+        num_firstepoch_steps = num_firstepoch_steps,
+    )
+    return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+def custom_raise_fall_scheduler_with_warmup(optimizer, num_warmup_steps, num_training_steps, num_firstepoch_steps, last_epoch=-1):
+    """
+    Args:
+        optimizer ([`~torch.optim.Optimizer`]):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (`int`):
+            The number of steps for the warmup phase.
+        num_training_steps (`int`):
+            The total number of training steps.
+        last_epoch (`int`, *optional*, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+    """
+    
+    lr_lambda = partial(
+        _get_fp_cosine_raise_and_fall_lr_lambda,
+        num_warmup_steps=num_warmup_steps,
+        num_training_steps=num_training_steps,
+        num_firstepoch_steps = num_firstepoch_steps,
+    )
+    return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+def neftune_forward(self, input: torch.Tensor):
+    """
+    Implements the NEFTune forward pass for the model. Note this works only for
+    torch.nn.Embedding layers. This method is slightly adapted from the original source code
+    that can be found here: https://github.com/neelsjain/NEFTune
+
+    Args:
+        input (`torch.Tensor`):
+            The input tensor to the model.
+        noise_alpha (`float`):
+            The noise alpha value to use for the NEFTune forward pass.
+    """
+    embeddings = torch.nn.functional.embedding(
+        input, self.weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse
+    )
+
+    if self.training:
+        # Add noise to the embeddings
+        dims = torch.tensor(embeddings.size(1) * embeddings.size(2))
+        mag_norm = self.neftune_noise_alpha / torch.sqrt(dims)
+        embeddings = embeddings + torch.zeros_like(embeddings).uniform_(-mag_norm, mag_norm)
+
+    return embeddings    
+
+
+class FPNEFtuneTrainer(transformers.Trainer):
+    def __init__(self,neftune_noise_alpha:float = 0.0, model = None, *args, **kwargs):
+        self.neftune_noise_alpha = neftune_noise_alpha
+        if self.neftune_noise_alpha > 0.0:
+            model = self._activate_neftune(model)
+        super().__init__(model = model, *args, **kwargs)
+
+   
+    def _activate_neftune(self, model):
+        r"""
+        Activates the neftune as presented in this code: https://github.com/neelsjain/NEFTune and paper: https://arxiv.org/abs/2310.05914
+        """
+        print(f"Activating {RED}NEFtune{RESET} with scale: {self.neftune_noise_alpha}")
+        if isinstance(model, transformers.PreTrainedModel):
+            embeddings = model.get_input_embeddings()
+        elif isinstance(model, PeftModel):
+            embeddings = model.base_model.get_input_embeddings()
+
+        embeddings.neftune_noise_alpha = self.neftune_noise_alpha
+        old_forward = embeddings.forward
+
+        # This hack seems to be needed to properly use a custom forward pass
+        # all credits to: https://discuss.pytorch.org/t/how-can-i-replace-the-forward-method-of-a-predefined-torchvision-model-with-my-customized-forward-function/54224/11
+        bound_method = neftune_forward.__get__(embeddings, embeddings.__class__)
+        setattr(embeddings, "forward", bound_method)
+
+        # embeddings.forward = neftune_forward
+        embeddings._trl_old_forward = old_forward
+
+        return model
+    
+    def train(self, *args, **kwargs):
+        output = super().train(*args, **kwargs)
+
+        # After training we make sure to retrieve back the original forward pass method
+        # for the embedding layer
+        if self.neftune_noise_alpha is not None:
+
+            if isinstance(self.model, transformers.PreTrainedModel):
+                embeddings = self.model.get_input_embeddings()
+            elif isinstance(self.model, PeftModel):
+                embeddings = self.model.base_model.get_input_embeddings()
+
+            if hasattr(embeddings, "_trl_old_forward"):
+                embeddings.forward = embeddings._trl_old_forward
+                del embeddings._trl_old_forward
+                del embeddings.neftune_noise_alpha
+
+        return output
+
+
+class FPSchedulerTrainer(transformers.Trainer):
+    def __init__(self,neftune_noise_alpha:float = 0.0, model = None, *args, **kwargs):
+        self.neftune_noise_alpha = neftune_noise_alpha
+        if self.neftune_noise_alpha > 0.0:
+            model = self._activate_neftune(model)
+        super().__init__(model = model, *args, **kwargs)
+
+   
+    def _activate_neftune(self, model):
+        r"""
+        Activates the neftune as presented in this code: https://github.com/neelsjain/NEFTune and paper: https://arxiv.org/abs/2310.05914
+        """
+        print(f"Activating {RED}NEFtune{RESET} with scale: {self.neftune_noise_alpha}")
+        if isinstance(model, transformers.PreTrainedModel):
+            embeddings = model.get_input_embeddings()
+        elif isinstance(model, PeftModel):
+            embeddings = model.base_model.get_input_embeddings()
+
+        embeddings.neftune_noise_alpha = self.neftune_noise_alpha
+        old_forward = embeddings.forward
+
+        # This hack seems to be needed to properly use a custom forward pass
+        # all credits to: https://discuss.pytorch.org/t/how-can-i-replace-the-forward-method-of-a-predefined-torchvision-model-with-my-customized-forward-function/54224/11
+        bound_method = neftune_forward.__get__(embeddings, embeddings.__class__)
+        setattr(embeddings, "forward", bound_method)
+
+        # embeddings.forward = neftune_forward
+        embeddings._trl_old_forward = old_forward
+
+        return model
+    
+    def train(self, *args, **kwargs):
+        output = super().train(*args, **kwargs)
+
+        # After training we make sure to retrieve back the original forward pass method
+        # for the embedding layer
+        if self.neftune_noise_alpha is not None:
+
+            if isinstance(self.model, transformers.PreTrainedModel):
+                embeddings = self.model.get_input_embeddings()
+            elif isinstance(self.model, PeftModel):
+                embeddings = self.model.base_model.get_input_embeddings()
+
+            if hasattr(embeddings, "_trl_old_forward"):
+                embeddings.forward = embeddings._trl_old_forward
+                del embeddings._trl_old_forward
+                del embeddings.neftune_noise_alpha
+
+        return output
+
+
+    def create_scheduler(self, num_training_steps: int, optimizer: torch.optim.Optimizer = None):
+        #Setup the scheduler. The optimizer of the trainer must have been set up either before this method is called or passed as an argument.
+        
+        num_train_epochs = self.args.num_train_epochs
+        num_warmup_steps=self.args.get_warmup_steps(num_training_steps)
+        num_firstepoch_steps = math.ceil(num_training_steps/num_train_epochs)
+        num_warmup_acc = num_warmup_steps*self.args.gradient_accumulation_steps 
+        num_firstepoch_steps_acc = num_firstepoch_steps*self.args.gradient_accumulation_steps
+        num_training_steps_acc = num_training_steps*self.args.gradient_accumulation_steps
+
+        custom_scheduler_params.update({'dynamic_scheduler_stop': False})
+ 
+        print (f"Warm-up steps aligned to Gradient accumulation ({self.args.gradient_accumulation_steps}) = {num_warmup_acc} actual warmup steps")
+        if self.args.lr_scheduler_type == 'cosine':
+            
+            num_warmup_acc_min = min(num_warmup_acc, num_firstepoch_steps_acc)
+
+            if num_warmup_acc>num_firstepoch_steps_acc:
+                print(f"\033[1;31;1mWARNING: The number of warmup steps is set too high! It will be clamped to 1 epoch, essentially going from warmup to annealing.\033[0;37;0m")
+                print (f"FP Scheduler Warmup: 0-[{num_warmup_acc_min}], Hold [{num_warmup_acc_min}]-{num_firstepoch_steps_acc}, Annealing {num_firstepoch_steps_acc}-{num_training_steps_acc}")
+            else:
+                print (f"FP Scheduler Warmup: 0-{num_warmup_acc_min}, Hold {num_warmup_acc_min}-{num_firstepoch_steps_acc}, Annealing {num_firstepoch_steps_acc}-{num_training_steps_acc}")
+
+            self.lr_scheduler = custom_cosine_scheduler_with_warmup(
+                    optimizer=self.optimizer if optimizer is None else optimizer,
+                    num_warmup_steps=num_warmup_steps,
+                    num_training_steps=num_training_steps, 
+                    num_firstepoch_steps = num_firstepoch_steps,
+                )
+            self._created_lr_scheduler = True
+            return self.lr_scheduler
+        elif self.args.lr_scheduler_type == 'constant':
+           
+            half_step_acc = num_training_steps_acc//2
+            num_warmup_acc_min = min(num_warmup_acc, half_step_acc)
+
+            if num_warmup_acc>half_step_acc:
+                print(f"\033[1;31;1mWARNING: The number of warmup steps is set too high! It will be clamped to half of all epochs, essentially going from warmup to annealing in the middle.\033[0;37;0m")
+                print (f"FP Scheduler Warmup: 0-[{num_warmup_acc_min}], Hold [{num_warmup_acc_min}]-{half_step_acc}, Annealing {half_step_acc}-{num_training_steps_acc}")
+            else:
+                print (f"FP Scheduler Warmup: 0-{num_warmup_acc_min}, Hold {num_warmup_acc_min}-{half_step_acc}, Annealing {half_step_acc}-{num_training_steps_acc}")
+
+            self.lr_scheduler = custom_half_scheduler_with_warmup(
+                    optimizer=self.optimizer if optimizer is None else optimizer,
+                    num_warmup_steps=num_warmup_steps,
+                    num_training_steps=num_training_steps, 
+                    num_firstepoch_steps = num_firstepoch_steps,
+                )
+            self._created_lr_scheduler = True
+            return self.lr_scheduler
+        elif self.args.lr_scheduler_type == 'constant_with_warmup':
+           
+            half_step_acc = num_training_steps_acc//2
+            
+            if num_warmup_steps>0:
+                print(f"Warmup doesn't apply to this scheduler [Raise-Fall]")
+
+            print (f"Scheduler Raise: 0-{half_step_acc}, Fall {half_step_acc}-{num_training_steps_acc}")
+
+            self.lr_scheduler = custom_raise_fall_scheduler_with_warmup(
+                    optimizer=self.optimizer if optimizer is None else optimizer,
+                    num_warmup_steps=num_warmup_steps,
+                    num_training_steps=num_training_steps, 
+                    num_firstepoch_steps = num_firstepoch_steps,
+                )
+            self._created_lr_scheduler = True
+            return self.lr_scheduler        
+        else:
+            return  super().create_scheduler(num_training_steps=num_training_steps, optimizer=optimizer)

extensions/Training_PRO/matplotgraph.py

@@ -0,0 +1,62 @@
+import os
+import json
+
+def create_graph(lora_path, lora_name):
+    try:
+        import matplotlib.pyplot as plt
+        from matplotlib.ticker import ScalarFormatter
+        
+        peft_model_path = f'{lora_path}/training_graph.json'
+        image_model_path = f'{lora_path}/training_graph.png'
+        # Check if the JSON file exists
+        if os.path.exists(peft_model_path):
+            # Load data from JSON file
+            with open(peft_model_path, 'r') as file:
+                data = json.load(file)
+            # Extract x, y1, and y2 values
+            x = [item['epoch'] for item in data]
+            y1 = [item['learning_rate'] for item in data]
+            y2 = [item['loss'] for item in data]
+
+            # Create the line chart
+            fig, ax1 = plt.subplots(figsize=(10, 6))
+        
+
+            # Plot y1 (learning rate) on the first y-axis
+            ax1.plot(x, y1, 'b-', label='Learning Rate')
+            ax1.set_xlabel('Epoch')
+            ax1.set_ylabel('Learning Rate', color='b')
+            ax1.tick_params('y', colors='b')
+
+            # Create a second y-axis
+            ax2 = ax1.twinx()
+
+            # Plot y2 (loss) on the second y-axis
+            ax2.plot(x, y2, 'r-', label='Loss')
+            ax2.set_ylabel('Loss', color='r')
+            ax2.tick_params('y', colors='r')
+
+            # Set the y-axis formatter to display numbers in scientific notation
+            ax1.yaxis.set_major_formatter(ScalarFormatter(useMathText=True))
+            ax1.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
+
+            # Add grid
+            ax1.grid(True)
+
+            # Combine the legends for both plots
+            lines, labels = ax1.get_legend_handles_labels()
+            lines2, labels2 = ax2.get_legend_handles_labels()
+            ax2.legend(lines + lines2, labels + labels2, loc='best')
+
+            # Set the title
+            plt.title(f'{lora_name} LR and Loss vs Epoch')
+
+            # Save the chart as an image
+            plt.savefig(image_model_path)
+
+            print(f"Graph saved in {image_model_path}")
+        else:
+            print(f"File 'training_graph.json' does not exist in the {lora_path}")
+      
+    except ImportError:
+        print("matplotlib is not installed. Please install matplotlib to create PNG graphs")

extensions/Training_PRO/script.py

@@ -0,0 +1,1308 @@
+import os
+
+os.environ["WANDB_MODE"] = "offline"
+# os.environ["WANDB_DISABLED"] = "true"
+
+import json
+import math
+import random
+import shutil
+import sys
+import threading
+import time
+import traceback
+from datetime import datetime
+from pathlib import Path
+
+import gradio as gr
+import pandas as pd
+import torch
+import transformers
+
+from functools import partial
+
+from .custom_scheduler import FPSchedulerTrainer, FPNEFtuneTrainer
+
+from .matplotgraph import create_graph
+from .train_utils import get_available_loras_local, precise_cut, sliding_block_cut, download_file_from_url
+
+from datasets import Dataset, load_dataset
+from peft import (
+    LoraConfig,
+    get_peft_model,
+    prepare_model_for_kbit_training,
+    set_peft_model_state_dict
+)
+from peft.utils.other import \
+    TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING as model_to_lora_modules
+from transformers.models.auto.modeling_auto import (
+    MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+)
+
+from modules import shared, utils
+from modules.ui import create_refresh_button
+
+from modules.evaluate import (
+    calculate_perplexity,
+    generate_markdown_table,
+    save_past_evaluations
+)
+from modules.logging_colors import logger
+from modules.models import reload_model
+from modules.utils import natural_keys
+
+import warnings
+warnings.filterwarnings(action = "ignore", message="torch.utils.checkpoint:")
+warnings.filterwarnings(action = "ignore", message="`do_sample` is set to `False`")
+
+params = {
+        "display_name": "Training PRO",
+        "is_tab": True
+}
+
+non_serialized_params = {
+        "debug_slicer": False,
+        "Lora_sortedByTime": False,
+        "stop_at_loss": 0,
+        "save_steps_under_loss": 0.0,
+        "save_checkpoint_now": False,
+        "training_loop": False,
+        "current_stability": 0,
+        "save_epochs": 0,
+        "checkpoint_offset": 0,
+        "epoch_offset":0,
+        "safe_serialization": False,
+}
+
+MODEL_CLASSES = {v[1]: v[0] for v in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.items()}
+
+PARAMETERS = ["lora_name", "always_override", "save_steps", "micro_batch_size", "batch_size", "epochs", "learning_rate", "lr_scheduler_type", "lora_rank", "lora_alpha", "lora_dropout", "cutoff_len", "dataset", "eval_dataset", "format", "eval_steps", "raw_text_file", "higher_rank_limit", "warmup_steps", "optimizer", "hard_cut_string", "train_only_after", "stop_at_loss", "add_eos_token", "min_chars", "report_to", "precize_slicing_overlap", "add_eos_token_type", "save_steps_under_loss", "add_bos_token", "training_projection","sliding_window","warmup_ratio","grad_accumulation","neft_noise_alpha"]
+WANT_INTERRUPT = False
+
+train_log = {}
+train_template = {}
+train_log_graph = []
+train_choices = ["all","q-k-v-o","q-k-v","k-v-down","q-v"]
+
+statistics = {
+			'loss': [],
+			'lr': [],
+}
+
+RED = "\033[91m"
+YELLOW = "\033[93m"
+GREEN = "\033[92m"
+RESET = "\033[0m"
+
+def ui():
+
+    with gr.Tab('Train LoRA', elem_id='lora-train-tab'):
+        tmp = gr.State('')
+        with gr.Row():
+            with gr.Column():
+                # YY.MM.DD
+                gr.Markdown("`Ver: 23.10.20 (REV2)` This is enhanced version of QLora Training. [Maintained by FP](https://github.com/FartyPants/Training_PRO/tree/main)")
+
+                with gr.Row():
+                    with gr.Column(scale=5):
+                        with gr.Row():
+                            copy_from = gr.Dropdown(label='Copy parameters from', value='None', choices=get_available_loras_local(non_serialized_params['Lora_sortedByTime']), elem_classes=['slim-dropdown'])
+                            create_refresh_button(copy_from, lambda: None, lambda: {'choices': get_available_loras_local(non_serialized_params['Lora_sortedByTime'])}, 'refresh-button')
+                    with gr.Column():
+                        sort_byTime = gr.Checkbox(label='Sort list by Date', value=False, info='Sorts Loras by date created.', elem_classes=['no-background'])                        
+
+                with gr.Row():
+                    with gr.Column(scale=5):
+                        lora_name = gr.Textbox(label='Name', info='The name of your new LoRA file')
+    
+                    with gr.Column():
+                        always_override = gr.Checkbox(label='Override Existing Files', value=False, info='If the name is the same, checking will replace the existing file, and unchecking will load and continue from it (the rank must be the same).', elem_classes=['no-background'])
+
+                with gr.Row():
+                    with gr.Column():
+                        lora_rank = gr.Slider(label='LoRA Rank', value=32, minimum=0, maximum=1024, step=4, info='Also called dimension count. Higher values = larger file, more content control. Smaller values = smaller file, less control. Use 4 or 8 for style, 128 or 256 to teach, 1024+ for fine-detail on big data. More VRAM is needed for higher ranks.')
+                        lora_alpha = gr.Slider(label='LoRA Alpha', value=64, minimum=0, maximum=2048, step=4, info='This divided by the rank becomes the scaling of the LoRA. Higher means stronger. A good standard value is twice your Rank.')
+                        batch_size = gr.Slider(visible= False, label='Batch Size', value=0, minimum=0, maximum=1024, step=4, info='Now Replaced with Gradient accumulation. Keeping it for sake of old saved data')
+                        micro_batch_size = gr.Slider(label='True Batch Size', value=4, minimum=1, maximum=128, step=1, info='Specifies how many text blocks per step will be trained. The higher value, the better the concept of training will be, but it requires more GPU memory and it reduces speed.')
+                        grad_accumulation = gr.Slider(label='Gradient Accumulation Steps', value=1, minimum=1, maximum=256, step=1, info="Virtually multiplies the Batch Size by averaging the learning over more than one step. VRAM friendly. Evens out loss fluctuations but can also degrade training fidelity.")
+
+                    with gr.Column():
+                        stop_at_loss = gr.Slider(label='Stop at loss (Can be changed during training)', minimum=0.0, maximum=3.0, step=0.1, value=0.00, info='The process will automatically stop once the desired loss value is reached.')
+                        gr.Markdown(" ")
+                        epochs = gr.Number(label='Epochs', value=3, info='Number of times every entry in the dataset should be fed into training. So 1 means feed each item in once, 5 means feed it in five times, etc.')
+                        learning_rate = gr.Textbox(label='Learning Rate', value='3e-4', info='In scientific notation. 3e-4 is a good starting base point. 1e-2 is extremely high, 1e-6 is extremely low.')
+                        lr_scheduler_type = gr.Dropdown(label='LR Scheduler', value='linear', choices=['linear', 'constant', 'constant_with_warmup', 'cosine', 'cosine_with_restarts', 'polynomial', 'inverse_sqrt', 'FP_low_epoch_annealing', 'FP_half_time_annealing','FP_raise_fall_creative'], info='Learning rate scheduler - defines how the learning rate changes over time. Custom schedulers: FP_low_epoch_annealing, FP_half_time_annealing, FP_raise_fall_creative (see README)', elem_classes=['slim-dropdown'])
+                        
+                with gr.Accordion(label='Checkpoints', open=True):
+                    with gr.Row():
+                        with gr.Column():
+                            save_steps = gr.Number(label='Save every n steps', value=0, info='A checkpoint will be saved every n steps and at each Epoch boundary. (0 = OFF)')
+                        with gr.Column():    
+                            save_steps_under_loss = gr.Slider(label='Save at 10% Loss change', value=1.8, minimum=0.0, maximum=3.0, step=0.1, info="Saves checkpoints at (or bellow) this loss and then each time loss falls by at least 10% This works independently from 'Save every n steps'")    
+                    with gr.Row():        
+                        save_chackpoint_now = gr.Button('Queue Checkpoint Now')
+
+                with gr.Accordion(label='Advanced Options', open=True):
+                    with gr.Row():
+                        with gr.Column():
+                            warmup_steps = gr.Number(label='Warmup Steps', value=100, info='Number of max steps used for a linear warmup. Reduces early over-fitting by the first training blocks. Value has precedent over Warmup Ratio. Aligns to the closest multiple of graddient accumulation')
+                            warmup_ratio = gr.Slider(label='Warmup Ratio', minimum=0.0, maximum=0.2, step=0.025, value=0.0, info='Ratio of total training steps that will be used for a linear warmup. It applies only if Warmup Step is 0.')
+                            neft_noise_alpha = gr.Slider(label='NEFtune noise scale', minimum=0.0, maximum=15, step=1, value=0.0, info='Add noise to the training to improve generalization. [0 - OFF, Starting value to experiment: 5]')
+                            training_projection = gr.Radio(value = train_choices[4], label='LLaMA Target Projections', info='Change the targets (LORA is typically q-v)', choices=train_choices)    
+                            lora_dropout = gr.Slider(label='LoRA Dropout', minimum=0.0, maximum=1.0, step=0.025, value=0.05, info='Percentage probability for dropout of LoRA layers. This can help reduce overfitting. Most users should leave at default.')
+                            optimizer = gr.Dropdown(label='Optimizer', value='adamw_torch', choices=['adamw_hf', 'adamw_torch', 'adamw_torch_fused', 'adamw_torch_xla', 'adamw_apex_fused', 'adafactor', 'adamw_bnb_8bit', 'adamw_anyprecision', 'sgd', 'adagrad'], info='Different optimizer implementation options, for advanced users. Effects of different options are not well documented yet.', elem_classes=['slim-dropdown'])
+
+                        with gr.Column():
+                            train_only_after = gr.Textbox(label='Train Only After', value='', info='Only consider text *after* this string in any given chunk for training. For Alpaca datasets, use "### Response:" to only train the response and ignore the input.')
+                            add_bos_token = gr.Checkbox(label='Add BOS token', value=True, info="Adds BOS token for each dataset item")
+                            add_eos_token = gr.Checkbox(label='Add EOS token', value=False, info="Adds EOS token for each dataset item")
+                            add_eos_token_type = gr.Dropdown(label='EOS placement (Text file)', choices=['Every Block', 'Hard Cut Blocks Only'], value='Every Block', info='', allow_custom_value = False)
+                            
+                            higher_rank_limit = gr.Checkbox(label='Enable higher ranks', value=False, info='If checked, changes Rank/Alpha slider above to go much higher. This will not work without a datacenter-class GPU.')
+                            report_to = gr.Radio(label="Save detailed logs with", value="None", choices=["None", "wandb", "tensorboard"], interactive=True)
+                # for future            
+                #with gr.Accordion(label='Dynamic Scheduler', open = False):
+                #    ds_min_epochs = gr.Number(label='Minimum Epochs', value='1', info='Minimum epochs that will be always performed before ramp down can be triggered')
+                #    ds_max_epochs = gr.Number(label='Maximum Epochs (fallback)', value='50', info='Maximum Epochs before the training will bail out completely (should be a large number)')
+                #    ds_loss_trigger = gr.Slider(label='Trigger Loss', minimum=0.0, maximum=2.8, step=0.1, value=1.6, info='Loss at which the ramp down schedule will be triggered')
+                #    ds_loss_rolling_window = gr.Number(label='Loss rolling average', value='4', info='Calculate loss by averaging last x numbers to avoid jumps and noise')
+                #    ds_epochs_to_ramp = gr.Slider(label='Ramp down ratio', minimum=0.0, maximum=2.0, step=0.1, value=1.00, info='How long the ramp down will last relative to ellapsed steps (before trigger)')
+                #    gr.Markdown('These are settings for FP_dynamic_loss_trigger scheduler. The scheduler will do warm up, then hold constant untill a loss falls under Trigger Loss, then it will commence linear ramp down schedule and stop. The length of ramp down is set by Ramp down ratio where (ramp down steps) = ratio * (elapsed steps). (The time to completition shown will be very high untill ramp down is triggered.)')
+                        
+
+            with gr.Column():
+                with gr.Tab(label='Formatted Dataset'):
+                    with gr.Row():
+                        with gr.Column():
+                            with gr.Row():
+                                dataset = gr.Dropdown(choices=get_datasets('training/datasets', 'json'), value='None', label='Dataset', info='The dataset file to use for training.', elem_classes=['slim-dropdown'])
+                                create_refresh_button(dataset, lambda: None, lambda: {'choices': get_datasets('training/datasets', 'json')}, 'refresh-button')
+                            with gr.Row():
+                                eval_dataset = gr.Dropdown(choices=get_datasets('training/datasets', 'json'), value='None', label='Evaluation Dataset', info='The (optional) dataset file used to evaluate the model after training.', elem_classes=['slim-dropdown'])
+                                create_refresh_button(eval_dataset, lambda: None, lambda: {'choices': get_datasets('training/datasets', 'json')}, 'refresh-button')
+
+                        with gr.Column():
+                            with gr.Row():
+                                format = gr.Dropdown(choices=get_datasets('training/formats', 'json'), value='None', label='Data Format', info='The format file used to decide how to format the dataset input.', elem_classes=['slim-dropdown'])
+                                create_refresh_button(format, lambda: None, lambda: {'choices': get_datasets('training/formats', 'json')}, 'refresh-button')
+                            with gr.Row():
+                                eval_steps = gr.Number(label='Evaluate every n steps', value=100, info='If an evaluation dataset is given, test it every time this many steps pass.')
+
+                with gr.Tab(label="Text file"):
+                    with gr.Row():
+                        raw_text_file = gr.Dropdown(choices=get_datasets('training/datasets', 'txt'), value='None', label='Text file', info='The text file to use for training.', elem_classes=['slim-dropdown'])
+                        create_refresh_button(raw_text_file, lambda: None, lambda: {'choices': get_datasets('training/datasets', 'txt')}, 'refresh-button')
+
+                    with gr.Row():
+                        with gr.Column():
+                            precize_slicing_overlap = gr.Checkbox(label='Add Overlapping blocks', value = True)
+                            sliding_window = gr.Checkbox(label='DEMENTOR Long-form Learning by FP (Highly Experimental, use low epochs)', value = False, info='Deep Memorization Enforcement Through Overlapping and Repetition. (I named it, so shush). Special process for learning long-form text using low amount of epochs.')
+                            #debug_slicer = gr.Checkbox(label='Dump sentencelist.json to logs', value = non_serialized_params['debug_slicer'], info='Debug Slicer')
+
+                        with gr.Column():
+                            hard_cut_string = gr.Textbox(label='Hard Cut String', value='\\n\\n\\n', info='String that indicates a cut between logical blocks of text (ex. Ideas or Chapters). Helps prevent unwanted overlap between unrelated ideas.')
+                            min_chars = gr.Number(label='Ignore small blocks', value=0, info='Ignore Text blocks that have less or equal characters than this number.')
+                with gr.Tab(label="URL"):
+                    with gr.Row():
+                        with gr.Column():
+                            download_file_url = gr.Textbox(label='Download JSON or txt file to datasets (or formats) folder', value='',info='The URL of a file to download. If on github, make sure you get url of the raw file (https://raw.githubusercontent.com/...). If huggin face, make sure the url has /resolve/ in it not /blob/')
+                            with gr.Row():
+                                download_check_overwrite = gr.Checkbox(label='Overwrite', value=False, info='Overwrite if file exist')
+                                download_folder = gr.Radio(label="Destination", value='training/datasets', choices=['training/datasets', 'training/formats'], interactive=True)
+                            download_button = gr.Button('Download')
+                            download_status = gr.Textbox(label='Download Status', value='', interactive=False)
+                with gr.Row():
+                    with gr.Column():
+                        with gr.Row():
+                            cutoff_len = gr.Slider(label='Chunk Length (Cutoff Length)', minimum=32, maximum=2048, value=256, step=32, info='The maximum length of a chunk (in tokens). Applies to both JSON dataset and text files. Higher values require much more VRAM.')
+                with gr.Row():
+                    with gr.Column():
+                        check_dataset_btn = gr.Button('Verify Dataset/Text File and suggest data entries')    
+                        check_dataset_txt = gr.Textbox(label='Dataset info', value='')
+
+                with gr.Row():
+                    start_button = gr.Button("Start LoRA Training", variant='primary')
+                    stop_button = gr.Button("Interrupt")
+
+                with gr.Accordion(label="Graph", open=True):
+                    with gr.Row():
+                        # show_actions_button = False - we use old gradio
+                        plot_graph = gr.LinePlot(x="epoch", y="value", title="Loss Metrics", overlay_point=True, tooltip=["epoch", "value"], x_lim=[0, 1], y_lim=[0, 3.5], width=500, height=250) 
+ 
+                output = gr.Markdown(value="Ready")
+
+    with gr.Tab('Perplexity evaluation', elem_id='evaluate-tab'):
+        with gr.Row():
+            with gr.Column():
+                models = gr.Dropdown(utils.get_available_models(), label='Models', multiselect=True)
+                evaluate_text_file = gr.Dropdown(choices=['wikitext', 'ptb', 'ptb_new'] + get_datasets('training/datasets', 'txt')[1:], value='wikitext', label='Input dataset', info='The text file on which the model will be evaluated. The first options are automatically downloaded: wikitext, ptb, and ptb_new. The next options are your local text files under training/datasets.')
+                with gr.Row():
+                    with gr.Column():
+                        stride_length = gr.Slider(label='Stride', minimum=1, maximum=2048, value=512, step=1, info='Used to make the evaluation faster at the cost of accuracy. 1 = slowest but most accurate. 512 is a common value.')
+
+                    with gr.Column():
+                        max_length = gr.Slider(label='max_length', minimum=0, maximum=shared.settings['truncation_length_max'], value=0, step=1, info='The context for each evaluation. If set to 0, the maximum context length for the model will be used.')
+
+                with gr.Row():
+                    start_current_evaluation = gr.Button("Evaluate loaded model")
+                    start_evaluation = gr.Button("Evaluate selected models")
+                    stop_evaluation = gr.Button("Interrupt")
+
+            with gr.Column():
+                evaluation_log = gr.Markdown(value='')
+
+        evaluation_table = gr.Dataframe(value=generate_markdown_table(), interactive=True)
+        with gr.Row():
+            save_comments = gr.Button('Save comments', elem_classes="small-button")
+            refresh_table = gr.Button('Refresh the table', elem_classes="small-button")
+
+    # Training events
+    all_params = [lora_name, always_override, save_steps, micro_batch_size, batch_size, epochs, learning_rate, lr_scheduler_type, lora_rank, lora_alpha, lora_dropout, cutoff_len, dataset, eval_dataset, format, eval_steps, raw_text_file, higher_rank_limit, warmup_steps, optimizer, hard_cut_string, train_only_after, stop_at_loss, add_eos_token, min_chars, report_to, precize_slicing_overlap, add_eos_token_type, save_steps_under_loss, add_bos_token, training_projection,sliding_window,warmup_ratio,grad_accumulation, neft_noise_alpha]
+
+    def fix_old_version(batch_size_val,micro_batch_size_val, grad_accumulation_val):
+        if batch_size_val>0:
+            gradient_acc =  batch_size_val // micro_batch_size_val
+            print(f"Using Old version of Batch Size ({batch_size_val}) to set Gradient Accumulation: {gradient_acc}")
+            return gradient_acc
+
+        return grad_accumulation_val
+
+    
+    copy_from.change(partial(do_copy_params, all_params= all_params), copy_from, all_params).then(fix_old_version,[batch_size,micro_batch_size, grad_accumulation],grad_accumulation)
+    start_button.click(do_train, all_params, [output,plot_graph])
+    stop_button.click(do_interrupt, None, None, queue=False)
+    higher_rank_limit.change(change_rank_limit, [higher_rank_limit], [lora_rank, lora_alpha])
+
+    def trigger_stop_at_loss(stop_at_loss_value):
+        non_serialized_params.update({"stop_at_loss": stop_at_loss_value})
+        if non_serialized_params['training_loop']:
+            print(f"Queue: [Stop at loss Change] to {stop_at_loss_value}")
+
+
+    stop_at_loss.change(trigger_stop_at_loss, stop_at_loss, None)
+
+    def trigger_save_checkpoint():
+        non_serialized_params.update({"save_checkpoint_now": True})
+        if non_serialized_params['training_loop']:
+            print("Queue: [Save checkpoint] Checkpoint will be saved after the current step is finished.")
+        else:
+            print("Use during the training to save the checkpoint at any time.")
+
+
+    def update_button():
+        return gr.Button.update('[Checkpoint in Queue]', variant='stop', interactive=True)
+
+    def update_button2():
+        time.sleep(1.0)
+        return gr.Button.update('Queue Checkpoint Now', variant='secondary',interactive = True)
+
+    save_chackpoint_now.click(trigger_save_checkpoint, None, None).then(update_button, None,save_chackpoint_now).then(update_button2, None,save_chackpoint_now)
+
+    dataset_calc_params = [save_steps,micro_batch_size, epochs, cutoff_len, dataset, format, raw_text_file, warmup_steps, hard_cut_string, min_chars, precize_slicing_overlap,sliding_window,warmup_ratio,grad_accumulation]
+
+    def check_dataset(save_steps:int, micro_batch_size: int, epochs: int, cutoff_len: int, dataset:str, format:str, raw_text_file:str, warmup_steps:int, hard_cut_string:str, min_chars:int, precize_slicing_overlap:bool,sliding_window:bool,warmup_ratio:float,grad_accumulation:int):
+        result = "Specify JSON dastaset or Text file"
+        total_blocks = 0
+        if shared.tokenizer is None:
+            yield "Tokenizer is not available. Please Load some Model first."
+            return
+        
+        
+        if raw_text_file not in ['None', '']:
+            logger.info("Loading Text file...")
+            fullpath = clean_path('training/datasets', f'{raw_text_file}')
+            fullpath = Path(fullpath)
+            if fullpath.is_dir():
+                logger.info('Training path directory {}'.format(raw_text_file))
+                raw_text = ""
+                file_paths = sorted(fullpath.glob('*.txt'), key=lambda path: natural_keys(path.name))
+                for file_path in file_paths:
+                    if file_path.is_file():
+                        with file_path.open('r', encoding='utf-8') as file:
+                            raw_text += file.read().replace('\r', '')
+
+                        logger.info(f"Loaded training file: {file_path.name}")
+            else:
+                try:
+                    with open(clean_path('training/datasets', f'{raw_text_file}.txt'), 'r', encoding='utf-8') as file:
+                        raw_text = file.read().replace('\r', '')
+                except:
+                    yield f"{raw_text_file}.txt doesn't seem to exsist anymore... check your training/datasets folder"
+                    return
+            
+ 
+            if min_chars<0:
+                min_chars = 0
+
+            # == New more precise slicing on sentence boundary ==
+            if sliding_window:
+                text_chunks = sliding_block_cut(raw_text, min_chars, False, cutoff_len, hard_cut_string,non_serialized_params['debug_slicer'])
+            else:
+                text_chunks = precise_cut(raw_text, precize_slicing_overlap, min_chars, False, cutoff_len, hard_cut_string,non_serialized_params['debug_slicer'])
+
+            total_blocks = len(text_chunks)
+            result = f"Text: ({raw_text_file}.txt) has {total_blocks} blocks (Block Size {cutoff_len} tokens)"
+            del text_chunks
+       
+        else:
+            if dataset in ['None', '']:
+                yield "Select dataset or text file."
+                return 
+
+            if format in ['None', '']:
+                yield "Select format choice for dataset."
+                return
+
+            with open(clean_path('training/formats', f'{format}.json'), 'r', encoding='utf-8-sig') as formatFile:
+                format_data: dict[str, str] = json.load(formatFile)
+
+            def generate_prompt(data_point: dict[str, str]):
+                for options, data in format_data.items():
+                    if set(options.split(',')) == set(x[0] for x in data_point.items() if (type(x[1]) is str and len(x[1].strip()) > 0)):
+                        for key, val in data_point.items():
+                            if type(val) is str:
+                                data = data.replace(f'%{key}%', val)
+                        return data
+                raise RuntimeError(f'Data-point "{data_point}" has no keyset match within format "{list(format_data.keys())}"')
+
+            def tokenize_dummy(prompt):
+
+                input_ids = shared.tokenizer.encode(prompt, truncation=True, max_length=cutoff_len)
+                labels = [1] * len(input_ids)
+                input_ids = torch.tensor(input_ids)
+                return {
+                    "input_ids": input_ids,
+                    "labels": labels,
+                    "attention_mask": input_ids.ne(shared.tokenizer.pad_token_id),
+                }
+
+            def generate_and_tokenize_prompt(data_point):
+                prompt = generate_prompt(data_point)
+                return tokenize_dummy(prompt)
+
+            logger.info("Loading JSON datasets...")
+            data = load_dataset("json", data_files=clean_path('training/datasets', f'{dataset}.json'))
+            
+            data_keys = [] 
+
+            if data:
+                if 'train' in data:  # Check if the 'train' split exists in the dataset
+                    data_keys = list(data['train'][0].keys())
+                    print("Data Keys:", data_keys)
+            else:
+                print("The dataset is empty.")
+
+            train_data = data['train'].map(generate_and_tokenize_prompt, new_fingerprint='%030x' % random.randrange(16**30))
+            total_blocks = train_data.num_rows
+
+            result = f"Dataset: ({dataset}.json) has {total_blocks} blocks @ length = {cutoff_len} tokens\n(Keys: {data_keys} - Format: {format}.json): "
+
+            #for options, data in format_data.items():
+            #    format_keys = options.split(',')
+            #    result += f"{format_keys}, "
+            #result = result.rstrip()    
+            #result = result.rstrip(',')  
+
+        if total_blocks>0:
+            number_ofSteps = int(math.ceil(total_blocks / micro_batch_size) * epochs) 
+            num_stepsPer_epoch = int(math.ceil(number_ofSteps/epochs))
+            min_warm = math.ceil(100 / grad_accumulation)
+
+            warmup_steps_suggest = min(int(min_warm*grad_accumulation), int(math.ceil(number_ofSteps * 0.1)))
+            warmup_steps_suggest = min(warmup_steps_suggest,num_stepsPer_epoch)
+
+            save_each_n_min = int(math.ceil(number_ofSteps/10))
+            save_each_n_max = int(math.ceil(number_ofSteps/5))
+            gradient_accumulation_max = int(total_blocks)//micro_batch_size
+
+ 
+            result += f"\n[Batch Size: {micro_batch_size}, Epochs: {epochs}, Gradient Accumulation: {grad_accumulation}]\n"
+            result += f"Total number of steps: {number_ofSteps}\n"
+            result += f"Steps per each Epoch: {num_stepsPer_epoch}\n"
+            result += f"Suggestions:\n"
+            result += f"Checkpoints: Save every {save_each_n_min} - {save_each_n_max} steps (Current: {int(save_steps)})\n"
+            result += f"Warmup steps: {warmup_steps_suggest} (Current: {int(warmup_steps)})"
+            if gradient_accumulation_max < grad_accumulation: 
+                result += f"\n\nWARNING: Gradient Accumulation {grad_accumulation} is too high: It should be below {gradient_accumulation_max}"
+
+
+        yield result
+        return
+    
+    check_dataset_btn.click(check_dataset, dataset_calc_params ,check_dataset_txt)
+
+    # Evaluation events. For some reason, the interrupt event
+    # doesn't work with the .then() syntax, so I write them one
+    # by one in this ugly but functional way.
+    ev = start_evaluation.click(calculate_perplexity, [models, evaluate_text_file, stride_length, max_length], evaluation_log, show_progress=False)
+    start_evaluation.click(generate_markdown_table, None, evaluation_table, show_progress=False)
+
+    start_current_evaluation.click(lambda: ['current model'], None, tmp)
+    ev_cur = start_current_evaluation.click(calculate_perplexity, [tmp, evaluate_text_file, stride_length, max_length], evaluation_log, show_progress=False)
+    start_current_evaluation.click(generate_markdown_table, None, evaluation_table, show_progress=False)
+
+    stop_evaluation.click(None, None, None, cancels=[ev, ev_cur], queue=False)
+    refresh_table.click(generate_markdown_table, None, evaluation_table, show_progress=True)
+    save_comments.click(
+        save_past_evaluations, evaluation_table, None).then(
+        lambda: "Comments saved.", None, evaluation_log, show_progress=False)
+
+    def reload_lora():
+        return gr.Dropdown.update(choices=get_available_loras_local(non_serialized_params['Lora_sortedByTime']))
+ 
+    # nonserialized items
+
+    sort_byTime.change(lambda x: non_serialized_params.update({"Lora_sortedByTime": x}), sort_byTime, None).then(reload_lora,None,copy_from) 
+    #debug_slicer.change(lambda x: non_serialized_params.update({"debug_slicer": x}), debug_slicer, None)
+
+    def update_dataset():
+        return gr.update(choices=get_datasets('training/datasets', 'json')), gr.update(choices=get_datasets('training/datasets', 'txt'))
+
+    download_button.click(download_file_from_url, [download_file_url,download_check_overwrite,download_folder] , download_status).then(update_dataset,None,[dataset , raw_text_file])
+
+def get_datasets(path: str, ext: str):
+    # include subdirectories for raw txt files to allow training from a subdirectory of txt files
+    #if ext == "txt":
+    #    return ['None'] + sorted(set([k.stem for k in list(Path(path).glob('txt')) + list(Path(path).glob('*/')) if k.stem != 'put-trainer-datasets-here']), key=natural_keys)
+
+    return ['None'] + sorted(set([k.stem for k in Path(path).glob(f'*.{ext}') if k.stem != 'put-trainer-datasets-here']), key=natural_keys)
+
+def do_interrupt():
+    global WANT_INTERRUPT
+    WANT_INTERRUPT = True
+
+
+def do_copy_params(lora_name: str, all_params):
+
+    if lora_name:
+        f_name = f"{shared.args.lora_dir}/{clean_path(None, lora_name)}/training_parameters.json"
+        if Path(f_name).is_file():
+            with open(f_name, 'r', encoding='utf-8') as format_file:
+                params: dict[str, str] = json.load(format_file)
+        else:
+            params = {}
+    else:
+        params = {}        
+
+    result = list()
+    for i in range(0, len(PARAMETERS)):
+        key = PARAMETERS[i]
+        if key in params:
+            result.append(params[key])
+        else:
+            result.append(all_params[i])
+
+    return result
+
+
+def change_rank_limit(use_higher_ranks: bool):
+    mult = 2 if use_higher_ranks else 1
+    return {"maximum": 1024 * mult, "__type__": "update"}, {"maximum": 2048 * mult, "__type__": "update"}
+
+
+def clean_path(base_path: str, path: str):
+    """Strips unusual symbols and forcibly builds a path as relative to the intended directory."""
+    path = path.replace('\\', '/').replace('..', '_')
+    if base_path is None:
+        return path
+
+    return f'{Path(base_path).absolute()}/{path}'
+
+
+def backup_adapter(input_folder):
+    # Get the creation date of the file adapter_model.bin
+    try:
+        adapter_file = Path(f"{input_folder}/adapter_model.bin")
+        if adapter_file.is_file():
+
+            logger.info("Backing up existing LoRA adapter...")
+            creation_date = datetime.fromtimestamp(adapter_file.stat().st_ctime)
+            creation_date_str = creation_date.strftime("Backup-%Y-%m-%d")
+
+            # Create the new subfolder
+            subfolder_path = Path(f"{input_folder}/{creation_date_str}")
+            subfolder_path.mkdir(parents=True, exist_ok=True)
+
+            # Check if the file already exists in the subfolder
+            backup_adapter_file = Path(f"{input_folder}/{creation_date_str}/adapter_model.bin")
+            if backup_adapter_file.is_file():
+                print(" - Backup already exists. Skipping backup process.")
+                return
+
+            # Copy existing files to the new subfolder
+            existing_files = Path(input_folder).iterdir()
+            for file in existing_files:
+                if file.is_file():
+                    shutil.copy2(file, subfolder_path)
+    except Exception as e:
+        print("An error occurred in backup_adapter:", str(e))
+
+
+def calc_trainable_parameters(model):
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        num_params = param.numel()
+        # if using DS Zero 3 and the weights are initialized empty
+        if num_params == 0 and hasattr(param, "ds_numel"):
+            num_params = param.ds_numel
+
+        all_param += num_params
+        if param.requires_grad:
+            trainable_params += num_params
+
+    return trainable_params, all_param
+
+
+
+def do_train(lora_name: str, always_override: bool, save_steps: int, micro_batch_size: int, batch_size: int, epochs: int, learning_rate: str, lr_scheduler_type: str, lora_rank: int, lora_alpha: int, lora_dropout: float, cutoff_len: int, dataset: str, eval_dataset: str, format: str, eval_steps: int, raw_text_file: str, higher_rank_limit: bool, warmup_steps: int, optimizer: str, hard_cut_string: str, train_only_after: str, stop_at_loss: float, add_eos_token: bool, min_chars: int, report_to: str, precize_slicing_overlap: bool, add_eos_token_type: str, save_steps_under_loss: float, add_bos_token: bool, training_projection: str,sliding_window:bool,warmup_ratio:float, grad_accumulation: int,neft_noise_alpha:float):
+
+    if shared.args.monkey_patch:
+        from alpaca_lora_4bit.monkeypatch.peft_tuners_lora_monkey_patch import (
+            replace_peft_model_with_int4_lora_model
+        )
+        replace_peft_model_with_int4_lora_model()
+    
+    global train_log_graph
+    global WANT_INTERRUPT
+    WANT_INTERRUPT = False
+
+    statistics['loss'] = []
+
+    statistics['loss'].append({'epoch': 0, 'value': 0})
+    zero_pd = pd.DataFrame(statistics['loss'])
+
+    # == Input validation / processing ==
+    yield "Preparing the input...", zero_pd
+    lora_file_path = clean_path(None, lora_name)
+    if lora_file_path.strip() == '':
+        yield "Missing or invalid LoRA file name input.", zero_pd
+        return
+
+    lora_file_path = f"{Path(shared.args.lora_dir)}/{lora_file_path}"
+    actual_lr = float(learning_rate)
+    model_type = type(shared.model).__name__
+
+    if model_type in MODEL_CLASSES:
+        model_id = MODEL_CLASSES[model_type]
+    else:
+        model_id = "llama"
+        if model_type == "PeftModelForCausalLM":
+            if len(shared.lora_names) > 0:
+                yield "You are trying to train a LoRA while you already have another LoRA loaded. This will work, but may have unexpected effects. *(Will continue anyway in 5 seconds, press `Interrupt` to stop.)*", zero_pd
+                logger.warning("Training LoRA over top of another LoRA. May have unexpected effects.")
+            else:
+                yield "Model ID not matched due to LoRA loading. Consider reloading base model. *(Will continue anyway in 5 seconds, press `Interrupt` to stop.)*", zero_pd
+                logger.warning("Model ID not matched due to LoRA loading. Consider reloading base model.")
+        else:
+            yield "LoRA training has only currently been validated for LLaMA, OPT, GPT-J, and GPT-NeoX models. Unexpected errors may follow. *(Will continue anyway in 5 seconds, press `Interrupt` to stop.)*", zero_pd
+            logger.warning(f"LoRA training has only currently been validated for LLaMA, OPT, GPT-J, and GPT-NeoX models. (Found model type: {model_type})")
+
+        time.sleep(5)
+
+    if shared.args.loader == 'GPTQ-for-LLaMa' and not shared.args.monkey_patch:
+        yield "LoRA training with GPTQ-for-LLaMa requires loading with `--monkey-patch`", zero_pd
+        return
+
+    if cutoff_len <= 0 or micro_batch_size <= 0 or actual_lr <= 0 or lora_rank <= 0 or lora_alpha <= 0:
+        yield "Cannot input zeroes.", zero_pd
+        return
+
+    #in new version we dumped this in favor of grad_accumulation
+    #set it to zero fo new save
+    batch_size = 0
+
+    gradient_accumulation_steps = grad_accumulation #batch_size // micro_batch_size
+    shared.tokenizer.pad_token_id = 0
+    shared.tokenizer.padding_side = "left"
+
+    def encode(text, prepend_bos_token):
+       
+        result = shared.tokenizer.encode(text, truncation=True, max_length=cutoff_len)
+        # Check if the first two tokens are BOS
+        if len(result) >= 2 and result[:2] == [shared.tokenizer.bos_token_id, shared.tokenizer.bos_token_id]:
+            result = result[1:]
+
+        if not prepend_bos_token and result[0] == shared.tokenizer.bos_token_id:
+            result = result[1:]
+        return result
+
+    def tokenize(prompt, append_eos_token=False, prepend_bos_token = False):
+
+        if train_only_after == '' or train_only_after not in prompt:
+            input_ids = encode(prompt, prepend_bos_token)
+
+            if append_eos_token and input_ids[-1] != shared.tokenizer.eos_token_id and len(input_ids) < cutoff_len:
+                input_ids.append(shared.tokenizer.eos_token_id)
+
+            input_ids = [shared.tokenizer.pad_token_id] * (cutoff_len - len(input_ids)) + input_ids
+            
+            labels = [1] * len(input_ids)
+        else:
+            ind = prompt.index(train_only_after) + len(train_only_after)
+            before_tokens = encode(prompt[:ind], prepend_bos_token)
+            after_tokens = encode(prompt[ind:], False)
+
+            if append_eos_token and after_tokens[-1] != shared.tokenizer.eos_token_id:
+                after_tokens.append(shared.tokenizer.eos_token_id)
+
+            full_length = len(after_tokens) + len(before_tokens)
+            if full_length > cutoff_len:
+                after_tokens = after_tokens[:cutoff_len - len(before_tokens)]
+            else:
+                before_tokens = [shared.tokenizer.pad_token_id] * (cutoff_len - full_length) + before_tokens
+
+            input_ids = before_tokens + after_tokens
+            labels = [-100] * len(before_tokens) + [1] * len(after_tokens)
+
+        input_ids = torch.tensor(input_ids)
+        return {
+            "input_ids": input_ids,
+            "labels": labels,
+            "attention_mask": input_ids.ne(shared.tokenizer.pad_token_id),
+        }
+
+    train_template.clear()
+            
+    #reset stuff
+    print(f"*** LoRA: {lora_name} ***")
+    non_serialized_params.update({"stop_at_loss": stop_at_loss})
+    non_serialized_params.update({"save_steps_under_loss": save_steps_under_loss+0.01})
+    non_serialized_params.update({"save_checkpoint_now": False})
+    non_serialized_params.update({"training_loop": False})
+    non_serialized_params.update({"current_stability": 0})
+    non_serialized_params.update({"save_epochs": 0})
+    non_serialized_params.update({"checkpoint_offset": 0})
+    non_serialized_params.update({"epoch_offset": 0})
+    train_log_graph.clear()
+  
+     # == Prep the dataset, format, etc ==
+    if raw_text_file not in ['None', '']:
+        train_template["template_type"] = "raw_text"
+        logger.info("Loading text file...")
+        fullpath = clean_path('training/datasets', f'{raw_text_file}')
+        fullpath = Path(fullpath)
+        if fullpath.is_dir():
+            logger.info('Training path directory {}'.format(raw_text_file))
+            raw_text = ""
+            file_paths = sorted(fullpath.glob('*.txt'), key=lambda path: natural_keys(path.name))
+            for file_path in file_paths:
+                if file_path.is_file():
+                    with file_path.open('r', encoding='utf-8') as file:
+                        raw_text += file.read().replace('\r', '')
+
+                    logger.info(f"Loaded training file: {file_path.name}")
+        else:
+            with open(clean_path('training/datasets', f'{raw_text_file}.txt'), 'r', encoding='utf-8') as file:
+                raw_text = file.read().replace('\r', '')
+        
+        # FPHAM PRECISE SLICING        
+        if min_chars<0:
+            min_chars = 0
+
+        add_EOS_to_all = add_eos_token and add_eos_token_type == 'Every Block'
+        add_EOS_to_HC = add_eos_token and add_eos_token_type != 'Every Block'
+
+        #print (f"add_eos_token {add_eos_token}, add_EOS_to_all {add_EOS_to_all}, add_EOS_to_HC {add_EOS_to_HC}")
+
+        # == New more precise slicing on sentence boundary ==
+        if sliding_window:
+            text_chunks = sliding_block_cut(raw_text, min_chars, add_EOS_to_HC, cutoff_len, hard_cut_string,non_serialized_params['debug_slicer'])
+        else:
+            text_chunks = precise_cut(raw_text, precize_slicing_overlap, min_chars, add_EOS_to_HC, cutoff_len, hard_cut_string,non_serialized_params['debug_slicer'])
+
+        train_data = Dataset.from_list([tokenize(x, add_EOS_to_all, add_bos_token) for x in text_chunks])
+        if add_EOS_to_all:
+            print(f"Added EOS to {len(text_chunks)} blocks") 
+
+        print(f"All Data Blocks: {len(text_chunks)}")
+
+        del text_chunks
+        eval_data = None
+    else:
+        if dataset in ['None', '']:
+            yield "Missing dataset choice input, cannot continue.", zero_pd
+            return
+
+        if format in ['None', '']:
+            yield "Missing format choice input, cannot continue.", zero_pd
+            return
+
+        train_template["template_type"] = "dataset"
+
+        with open(clean_path('training/formats', f'{format}.json'), 'r', encoding='utf-8-sig') as formatFile:
+            format_data: dict[str, str] = json.load(formatFile)
+
+        # == store training prompt ==
+        for _, value in format_data.items():
+            prompt_key = f"template_{len(train_template)}"
+            train_template[prompt_key] = value
+
+        def generate_prompt(data_point: dict[str, str]):
+            for options, data in format_data.items():
+                if set(options.split(',')) == set(x[0] for x in data_point.items() if (type(x[1]) is str and len(x[1].strip()) > 0)):
+                    for key, val in data_point.items():
+                        if type(val) is str:
+                            data = data.replace(f'%{key}%', val)
+                    return data
+            raise RuntimeError(f'Data-point "{data_point}" has no keyset match within format "{list(format_data.keys())}"')
+
+        def generate_and_tokenize_prompt(data_point):
+            prompt = generate_prompt(data_point)
+            return tokenize(prompt, add_eos_token, add_bos_token)
+
+        logger.info("Loading JSON datasets...")
+        data = load_dataset("json", data_files=clean_path('training/datasets', f'{dataset}.json'))
+        train_data = data['train'].map(generate_and_tokenize_prompt, new_fingerprint='%030x' % random.randrange(16**30))
+
+        print(f"BOS: {add_bos_token} EOS: {add_eos_token}") 
+        print(f"Data Blocks: {train_data.num_rows}")
+
+        if eval_dataset == 'None':
+            eval_data = None
+        else:
+            eval_data = load_dataset("json", data_files=clean_path('training/datasets', f'{eval_dataset}.json'))
+            eval_data = eval_data['train'].map(generate_and_tokenize_prompt, new_fingerprint='%030x' % random.randrange(16**30))
+
+    # == We MUST reload model if it went through any previous training, even failed one ==
+    if shared.model_dirty_from_training:
+        selected_model = shared.model_name
+        if selected_model:
+            print("\033[1;31;1m(Model has been modified by previous training, it needs to be reloaded...)\033[0;37;0m")
+            try:
+                yield f"Reloading {selected_model}...", zero_pd
+                reload_model()
+                shared.tokenizer.pad_token_id = 0
+                shared.tokenizer.padding_side = "left"
+
+                if shared.model is not None:
+                    print("Model reloaded OK, continue with training.")
+                else:
+                    return f"Failed to load {selected_model}."
+            except:
+                exc = traceback.format_exc()
+                logger.error('Failed to reload the model.')
+                print(exc)
+                return exc.replace('\n', '\n\n')
+
+    # == Start prepping the model itself ==
+    if not hasattr(shared.model, 'lm_head') or hasattr(shared.model.lm_head, 'weight'):
+        logger.info("Getting model ready...")
+        # here we can disable gradient checkpoint, by default = true,  use_gradient_checkpointing=True
+        prepare_model_for_kbit_training(shared.model)
+
+    # base model is now frozen and should not be reused for any other LoRA training than this one
+    shared.model_dirty_from_training = True
+    print(f"Transformers Model Type: {YELLOW}{model_type}{RESET}")
+
+    if training_projection==train_choices[0]:
+        model_to_lora_modules[model_id] = ["gate_proj","down_proj","up_proj","q_proj","k_proj","v_proj","o_proj"]
+    elif training_projection==train_choices[1]:
+        model_to_lora_modules[model_id] = ["q_proj","k_proj", "v_proj", "o_proj"]
+    elif training_projection==train_choices[2]:
+        model_to_lora_modules[model_id] = ["q_proj","k_proj", "v_proj"]
+    elif training_projection==train_choices[3]:
+        model_to_lora_modules[model_id] = ["k_proj", "v_proj", "down_proj"]        
+    else:
+        model_to_lora_modules[model_id] = ["q_proj", "v_proj"]
+
+
+    logger.info("Preparing for training...")
+    config = LoraConfig(
+        r=lora_rank,
+        lora_alpha=lora_alpha,
+        target_modules=model_to_lora_modules[model_id],
+        lora_dropout=lora_dropout,
+        bias="none",
+        task_type="CAUSAL_LM"
+    )
+
+    # == Backup the existing adapter ==
+    if not always_override:
+        backup_adapter(lora_file_path)
+
+    # == get model trainable params
+    model_trainable_params, model_all_params = calc_trainable_parameters(shared.model)
+
+    try:
+        logger.info("Creating LoRA model...")
+        lora_model = get_peft_model(shared.model, config)
+        if not always_override and Path(f"{lora_file_path}/adapter_model.bin").is_file():
+            logger.info("Loading existing LoRA data...")
+            state_dict_peft = torch.load(f"{lora_file_path}/adapter_model.bin")
+            set_peft_model_state_dict(lora_model, state_dict_peft)
+
+            print(f" + Continue Training on {RED}{lora_file_path}/adapter_model.bin{RESET}")
+            
+            #load training_log.json if exist
+           
+            if Path(f"{lora_file_path}/training_log.json").is_file():
+                with open(f"{lora_file_path}/training_log.json", 'r') as json_file:
+                    json_ilog = json.load(json_file)
+                    for key, value in json_ilog.items():
+                        if key=='current_steps':
+                            non_serialized_params.update({"checkpoint_offset": int(value+1)})
+                            print(f" + Checkpoints will be saved with offset: {RED}{non_serialized_params['checkpoint_offset']}{RESET}")
+                        if key=='epoch':
+                            non_serialized_params.update({"epoch_offset": value})
+                            print(f" + Epoch offset: {RED}{non_serialized_params['epoch_offset']}{RESET}")
+           
+
+            if Path(f"{lora_file_path}/training_graph.json").is_file():
+                try:
+                    with open(f"{lora_file_path}/training_graph.json", 'r') as json_file:
+                        train_log_graph = json.load(json_file)
+                        print(" + Training Graph loaded")   
+                except:
+                    print(f"Can't read training_graph")
+
+
+    except:
+        yield traceback.format_exc().replace('\n', '\n\n'), zero_pd
+        return
+
+    if shared.args.monkey_patch:
+        from alpaca_lora_4bit.autograd_4bit import Autograd4bitQuantLinear
+        from alpaca_lora_4bit.models import Linear4bitLt
+        for _, m in lora_model.named_modules():
+            if isinstance(m, Autograd4bitQuantLinear) or isinstance(m, Linear4bitLt):
+                if m.is_v1_model:
+                    m.zeros = m.zeros.half()
+                m.scales = m.scales.half()
+
+    class Tracked():
+        def __init__(self):
+            self.current_steps = 0
+            self.max_steps = 0
+            self.did_save = False
+
+    tracked = Tracked()
+    actual_save_steps = math.ceil(save_steps / gradient_accumulation_steps)
+
+    class Callbacks(transformers.TrainerCallback):
+        def on_step_begin(self, args: transformers.TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, **kwargs):
+            tracked.current_steps = state.global_step * gradient_accumulation_steps
+            tracked.max_steps = state.max_steps * gradient_accumulation_steps
+            ssteps10 = int(max(2,(state.max_steps/epochs)*0.1))
+
+            if WANT_INTERRUPT:
+                control.should_epoch_stop = True
+                control.should_training_stop = True
+            else:
+                current_loss = float(train_log.get('loss', 0.0))
+                current_epoch_int = int(float(train_log.get('epoch', 0.0)))
+              
+                force_save = False
+
+                current_steps_offset = tracked.current_steps + non_serialized_params['checkpoint_offset']
+
+                folder_save = f"checkpoint-{current_steps_offset}"    
+
+                # save if triggered by user
+                if non_serialized_params['save_checkpoint_now']:
+                    force_save = True
+                    non_serialized_params.update({"save_checkpoint_now": False})
+                    print(f"\033[1;31;1mSave Checkpoint manually trigerred.\033[0;37;0m")
+                    folder_save = f"checkpoint-{current_steps_offset}-user"  
+
+                patience = 3     # Set the number of consecutive steps for tracking stability
+                
+                if gradient_accumulation_steps==1:
+                    patience = 4
+
+                min_steps = ssteps10
+
+                # Save each time the loss is below the threshold 
+                if current_loss < non_serialized_params['save_steps_under_loss'] and current_loss > 0 and state.global_step > min_steps:
+                    current_stability = non_serialized_params['current_stability']
+                    current_stability += 1
+                    non_serialized_params.update({"current_stability": current_stability}) 
+
+                    if current_stability >= patience:
+                        current_stability = 0
+                        non_serialized_params.update({"current_stability": current_stability})     
+                        current_loss_dec = round(current_loss, 2)
+                        loss_str = f"{current_loss_dec:.2f}"
+                        loss_str = loss_str.replace('.', '_')
+                        new_save = (current_loss_dec-0.1) + 0.01
+                        non_serialized_params.update({"save_steps_under_loss": new_save})
+
+                        folder_save = f"checkpoint-{current_steps_offset}-loss-{loss_str}" 
+                        force_save = True   
+
+                   
+                else:
+                    # Reset stability if the loss goes above the threshold
+                    non_serialized_params.update({"current_stability": 0})   
+
+                # Save full epochs
+                if actual_save_steps>0 and current_epoch_int > non_serialized_params['save_epochs'] and state.global_step > min_steps: 
+
+                    
+                    current_epoch_offset = current_epoch_int
+                    
+                    if non_serialized_params['epoch_offset'] > 0:
+                        current_epoch_offset = current_epoch_int + round(non_serialized_params['epoch_offset'], 2)
+                    
+                    ep_off_str = f"{current_epoch_offset}"
+                    ep_off_str = ep_off_str.replace('.', '_')
+                    folder_save = f"checkpoint-{current_steps_offset}-epoch-{ep_off_str}" 
+
+                    non_serialized_params.update({"save_epochs": current_epoch_int})
+                    force_save = True
+
+                # save each actual_save_steps
+                if state.global_step > 0 and actual_save_steps > 0 and state.global_step % actual_save_steps == 0:
+                    folder_save = f"checkpoint-{current_steps_offset}"  
+                    force_save = True   
+
+                if force_save:       
+                    lora_model.save_pretrained(f"{lora_file_path}/{folder_save}/", safe_serialization = non_serialized_params['safe_serialization'])
+                    print(f"\033[1;30;40mStep: {tracked.current_steps:6} \033[0;37;0m Saved: [{folder_save}]")
+                    # Save log
+                    with open(f"{lora_file_path}/{folder_save}/training_log.json", 'w', encoding='utf-8') as file:
+                        json.dump(train_log, file, indent=2)
+                    # == Save training prompt ==
+                    with open(f"{lora_file_path}/{folder_save}/training_prompt.json", 'w', encoding='utf-8') as file:
+                        json.dump(train_template, file, indent=2)
+                
+
+        def on_substep_end(self, args: transformers.TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, **kwargs):
+            tracked.current_steps += 1
+            if WANT_INTERRUPT:
+                control.should_epoch_stop = True
+                control.should_training_stop = True
+
+        def on_log(self, args: transformers.TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, logs, **kwargs):
+            train_log.update(logs)
+
+            current_steps_offset = tracked.current_steps + non_serialized_params['checkpoint_offset']
+            current_epoch_offset = train_log.get('epoch', 0.0) + non_serialized_params['epoch_offset']
+
+            train_log.update({"current_steps": tracked.current_steps})
+            train_log.update({"current_steps_adjusted": current_steps_offset})
+            train_log.update({"epoch_adjusted": current_epoch_offset})
+
+            if WANT_INTERRUPT:
+                print("\033[1;31;1mInterrupted by user\033[0;37;0m")
+
+            if non_serialized_params['checkpoint_offset']>0:
+                print(f"\033[1;30;40mStep: {tracked.current_steps:6} [+{non_serialized_params['checkpoint_offset']}] \033[0;37;0m", end='')
+            else:
+                print(f"\033[1;30;40mStep: {tracked.current_steps:6} \033[0;37;0m", end='')
+            
+            graphentry = {
+                'current_steps': int(train_log.get('current_steps_adjusted',0)),
+                'loss': float(train_log.get('loss', 0.0)),
+                'learning_rate': float(train_log.get('learning_rate', 0.0)),
+                'epoch': float(train_log.get('epoch_adjusted', 0.0))
+            }
+
+            cur_loss = float(train_log.get('loss', 0.0))
+            cur_lr = float(train_log.get('learning_rate', 0.0))
+            cur_epoch = float(train_log.get('epoch', 0.0))
+            
+            if len(statistics['loss']) == 1:
+                first_epoch = statistics['loss'][0]['epoch']
+                first_value = statistics['loss'][0]['value']
+                if first_value ==0:
+                     statistics['loss'] = []
+
+
+            statistics['loss'].append({'epoch': cur_epoch, 'value': cur_loss})
+            statistics['lr'].append({'epoch': cur_epoch, 'value': cur_lr})
+
+            # Add the entry to the continuous log
+            train_log_graph.append(graphentry)
+
+            # Save the graph log for now, we can later generate full graph
+            with open(f"{lora_file_path}/training_graph.json", 'w') as file:
+                json.dump(train_log_graph, file, indent=4)
+
+            if 'loss' in logs:
+                loss = float(logs['loss'])
+                if loss <= stop_at_loss:
+                    control.should_epoch_stop = True
+                    control.should_training_stop = True
+                    print(f"{RED}Stop Loss {stop_at_loss} reached.{RESET}")
+
+    # FPHAM SAMPLE REQ Transformers error handling
+    gradient_accumulation_max = int(train_data.num_rows)//micro_batch_size
+    
+    if gradient_accumulation_max < gradient_accumulation_steps:
+        print(f"{RED}WARNING:{RESET} Current gradient accumulation is {RED}too high{RESET} for the amount of training data.")
+        print(f"Gradient accumulation: {gradient_accumulation_steps} should be less than: {gradient_accumulation_max}. {RED}This could crash Accelerate/Transformers{RESET}")
+        #min_batchSize = sample_req*micro_batch_size
+        print(f"Preferable fix: {RED}Increase the size of dataset{RESET}")
+        print(f"... or Decrerase Gradient Accumulation {RED}{gradient_accumulation_steps}{RESET} to below {GREEN}{gradient_accumulation_max}{RESET}")
+        gradient_accumulation_steps = max(1,gradient_accumulation_max-1)
+        print(f"Last resort fix for this run: Lowering Gradient accumulation to {GREEN}{gradient_accumulation_steps}{RESET} [Good luck]")
+
+    else:
+        print(f"Data Size Check: Gradient accumulation: {YELLOW}{gradient_accumulation_steps}{RESET} <= Blocks/Batch {gradient_accumulation_max} ... {GREEN}[OK]{RESET}")
+
+    #END OF FPHAM SAMPLE REQ
+
+    # FPHAM Custom Scheduler ==
+    custom_scheduller = False
+    lr_scheduler_type_arg = lr_scheduler_type
+
+    if lr_scheduler_type == 'FP_low_epoch_annealing':
+        custom_scheduller = True
+        lr_scheduler_type_arg = 'cosine'
+    elif lr_scheduler_type == 'FP_half_time_annealing':
+        custom_scheduller = True
+        lr_scheduler_type_arg = 'constant'
+    elif lr_scheduler_type =='FP_raise_fall_creative':
+        custom_scheduller = True
+        lr_scheduler_type_arg = 'constant_with_warmup'
+    
+    #gradient_checkpointing=True
+    
+    args=transformers.TrainingArguments(
+            report_to=report_to if report_to != "None" else None,
+            per_device_train_batch_size=micro_batch_size,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            warmup_steps=math.ceil(warmup_steps / gradient_accumulation_steps),
+            warmup_ratio = warmup_ratio,
+            num_train_epochs=epochs,
+            learning_rate=actual_lr,
+            fp16=False if shared.args.cpu else True,
+            optim=optimizer,
+            logging_steps=1,
+            evaluation_strategy="steps" if eval_data is not None else "no",
+            eval_steps=math.ceil(eval_steps / gradient_accumulation_steps) if eval_data is not None else None,
+            save_strategy="steps" if eval_data is not None else "no",
+            output_dir=lora_file_path,
+            lr_scheduler_type=lr_scheduler_type_arg,
+            load_best_model_at_end=eval_data is not None,
+            # TODO: Enable multi-device support
+            ddp_find_unused_parameters=None,
+            no_cuda=shared.args.cpu,
+        )
+
+    if custom_scheduller:
+        trainer = FPSchedulerTrainer(
+            neftune_noise_alpha=neft_noise_alpha,
+            model=lora_model,
+            train_dataset=train_data,
+            eval_dataset=eval_data,
+            args=args,
+            data_collator=transformers.DataCollatorForLanguageModeling(shared.tokenizer, mlm=False),
+            callbacks=list([Callbacks()])
+        )
+    elif neft_noise_alpha > 0:
+            trainer = FPNEFtuneTrainer(
+            neftune_noise_alpha=neft_noise_alpha,
+            model=lora_model,
+            train_dataset=train_data,
+            eval_dataset=eval_data,
+            args=args,
+            data_collator=transformers.DataCollatorForLanguageModeling(shared.tokenizer, mlm=False),
+            callbacks=list([Callbacks()])
+        )
+    else:
+        trainer = transformers.Trainer(
+            model=lora_model,
+            train_dataset=train_data,
+            eval_dataset=eval_data,
+            args=args,
+            data_collator=transformers.DataCollatorForLanguageModeling(shared.tokenizer, mlm=False),
+            callbacks=list([Callbacks()])
+        )
+    
+    # END OF FPHAM CUSTOM SCHEDULER
+
+    lora_model.config.use_cache = False
+
+    if torch.__version__ >= "2" and sys.platform != "win32":
+        lora_model = torch.compile(lora_model)
+
+    # == Save parameters for reuse ==
+    with open(f"{lora_file_path}/training_parameters.json", 'w', encoding='utf-8') as file:
+        vars = locals()
+        json.dump({x: vars[x] for x in PARAMETERS}, file, indent=2)
+
+    # == Save training prompt ==
+    with open(f"{lora_file_path}/training_prompt.json", 'w', encoding='utf-8') as file:
+        json.dump(train_template, file, indent=2)
+
+    # == Main run and monitor loop ==
+    logger.info("Starting training...")
+    yield "Starting...", zero_pd
+
+    lora_trainable_param, lora_all_param = calc_trainable_parameters(lora_model)
+
+    projections_string = ", ".join([projection.replace("_proj", "") for projection in model_to_lora_modules[model_id]])
+
+    print(f"Training '{model_id}' model using {YELLOW}({projections_string}){RESET} projections")
+
+    if lora_all_param > 0:
+        print(f"Trainable params: {lora_trainable_param:,d} ({RED}{100 * lora_trainable_param / lora_all_param:.4f} %{RESET}), All params: {lora_all_param:,d} (Model: {model_all_params:,d})")
+
+    train_log.update({"base_model_name": shared.model_name})
+    train_log.update({"base_model_class": shared.model.__class__.__name__})
+    train_log.update({"base_loaded_in_4bit": getattr(lora_model, "is_loaded_in_4bit", False)})
+    train_log.update({"base_loaded_in_8bit": getattr(lora_model, "is_loaded_in_8bit", False)})
+    train_log.update({"projections": projections_string})
+    if non_serialized_params['checkpoint_offset'] > 0:
+        train_log.update({"last_run_steps_offset": non_serialized_params['checkpoint_offset']})
+        train_log.update({"last_run_epoch_offset": non_serialized_params['epoch_offset']})
+
+
+    if non_serialized_params['checkpoint_offset'] > 0:
+        print(f"Continue training on {RED}previous adapter{RESET} from epoch: {RED}{non_serialized_params['epoch_offset']}{RESET}")
+
+    if stop_at_loss > 0:
+        print(f"Monitoring loss {RED}(Auto-Stop at: {stop_at_loss}){RESET}")
+
+    
+
+    if WANT_INTERRUPT:
+        yield "Interrupted before start.", zero_pd
+        return
+
+    def log_train_dataset(trainer):
+        decoded_entries = []
+        # Try to decode the entries and write the log file
+        try:
+            # Iterate over the first 10 elements in the dataset (or fewer if there are less than 10)
+            for i in range(min(10, len(trainer.train_dataset))):
+                decoded_text = shared.tokenizer.decode(trainer.train_dataset[i]['input_ids'])
+                decoded_entries.append({"value": decoded_text})
+
+            # Write the log file
+            Path('logs').mkdir(exist_ok=True)
+            with open(Path('logs/train_dataset_sample.json'), 'w') as json_file:
+                json.dump(decoded_entries, json_file, indent=4)
+
+            logger.info("Log file 'train_dataset_sample.json' created in the 'logs' directory.")
+        except Exception as e:
+            logger.error(f"Failed to create log file due to error: {e}")
+
+    def threaded_run():
+        log_train_dataset(trainer)
+        trainer.train()
+        # Note: save in the thread in case the gradio thread breaks (eg browser closed)
+        lora_model.save_pretrained(lora_file_path, safe_serialization = non_serialized_params['safe_serialization'])
+        logger.info("LoRA training run is completed and saved.")
+        # Save log
+        with open(f"{lora_file_path}/training_log.json", 'w', encoding='utf-8') as file:
+            json.dump(train_log, file, indent=2)
+
+    thread = threading.Thread(target=threaded_run)
+    thread.start()
+    last_step = 0
+    start_time = time.perf_counter()
+
+    while thread.is_alive():
+        time.sleep(0.5)
+
+        if statistics['loss']:
+            max_value_dict = max(statistics['loss'], key=lambda x: x['value'])
+            max_value = max_value_dict['value']+0.4
+            first_epoch = statistics['loss'][0]['epoch']
+            last_epoch = statistics['loss'][-1]['epoch']
+        else:
+            max_value = 3.5
+            last_epoch = 0
+            first_epoch = 0           
+
+        if WANT_INTERRUPT:
+
+            losses = gr.LinePlot.update(
+				value = pd.DataFrame(statistics['loss']),
+                x="epoch", y="value",
+                title="Loss Metrics",
+                overlay_point=True, tooltip=["epoch", "value"],
+				x_lim=[first_epoch,last_epoch], y_lim=[0,max_value],
+                width=500, height=250 )
+
+            yield "Interrupting, please wait... *(Run will stop after the current training step completes.)*", losses
+
+        elif tracked.current_steps != last_step:
+            last_step = tracked.current_steps
+            time_elapsed = time.perf_counter() - start_time
+            lastloss = float(train_log.get('loss', 0.0))
+
+            non_serialized_params.update({"training_loop": True})               
+
+            if lastloss > 0:
+                lastloss_str = f", ... Current Loss: `{lastloss:.2f}`"
+            else:
+                lastloss_str = ""
+
+            if time_elapsed <= 0:
+                timer_info = ""
+                total_time_estimate = 999
+            else:
+                its = tracked.current_steps / time_elapsed
+                if its > 1:
+                    timer_info = f"`{its:.2f}` it/s"
+                else:
+                    timer_info = f"`{1.0/its:.2f}` s/it"
+
+                total_time_estimate = (1.0 / its) * (tracked.max_steps)
+
+            if stop_at_loss != non_serialized_params['stop_at_loss']:
+                stop_at_loss = non_serialized_params['stop_at_loss']
+                print(f"Stop at loss changed {RED}(Auto-Stop at: {stop_at_loss}){RESET}")
+            
+            losses = gr.LinePlot.update(
+				value = pd.DataFrame(statistics['loss']),
+                x="epoch", y="value",
+                title="Loss Metrics",
+                overlay_point=True, tooltip=["epoch", "value"],
+				x_lim=[first_epoch,last_epoch], y_lim=[0,max_value],
+                width=500, height=250 )
+				
+
+            yield f"Running... **{tracked.current_steps}** / **{tracked.max_steps}** ... {timer_info}, {format_time(time_elapsed)} / {format_time(total_time_estimate)} ... {format_time(total_time_estimate - time_elapsed)} remaining {lastloss_str}", losses
+
+    # Saving in the train thread might fail if an error occurs, so save here if so.
+
+    #return_pd = pd.DataFrame(statistics['loss'])
+
+    if statistics['loss']:
+        max_value_dict = max(statistics['loss'], key=lambda x: x['value'])
+        max_value = max_value_dict['value']+0.4
+        first_epoch = statistics['loss'][0]['epoch']
+        last_epoch = statistics['loss'][-1]['epoch']
+    else:
+        max_value = 3.5
+        last_epoch = 0
+        first_epoch = 0 
+
+    return_pd = gr.LinePlot.update(
+        value = pd.DataFrame(statistics['loss']),
+        x="epoch", y="value",
+        title="Loss Metrics",
+        overlay_point=True, tooltip=["epoch", "value"],
+        x_lim=[first_epoch,last_epoch], y_lim=[0,max_value],
+        width=500, height=250)
+
+    non_serialized_params.update({"training_loop": False})
+
+    if not tracked.did_save:
+        logger.info("Training complete, saving...")
+        lora_model.save_pretrained(lora_file_path, safe_serialization = non_serialized_params['safe_serialization'])
+
+    if WANT_INTERRUPT:
+        logger.info("Training interrupted.")
+        yield f"Interrupted by user. LoRA saved to `{lora_file_path}`.", return_pd
+    else:
+        logger.info("Training complete!")
+        yield f"Done! LoRA saved to `{lora_file_path}`.\n\nBefore testing your new LoRA, make sure to first reload the model, as it is currently dirty from training.", return_pd
+
+    create_graph(lora_file_path, lora_name)
+
+def format_time(seconds: float):
+    if seconds < 120:
+        return f"`{seconds:.0f}` seconds"
+
+    minutes = seconds / 60
+    if minutes < 120:
+        return f"`{minutes:.0f}` minutes"
+
+    hours = minutes / 60
+    return f"`{hours:.0f}` hours"

extensions/Training_PRO/train_utils.py

@@ -0,0 +1,368 @@
+import os
+from modules import shared, utils
+from pathlib import Path
+import requests
+import tqdm
+import json
+
+'''
+def get_gpu_memory_usage(rank):
+    return {
+        'total': round(torch.cuda.get_device_properties(rank).total_memory / (1024**3), 2),
+        'max': round(torch.cuda.max_memory_allocated(rank) / (1024**3), 2),
+        'reserved': round(torch.cuda.memory_reserved(rank) / (1024**3), 2),
+        'allocated': round(torch.cuda.memory_allocated(rank) / (1024**3), 2)
+    }
+'''
+
+def list_subfoldersByTime(directory):
+
+    if not directory.endswith('/'):
+        directory += '/'
+    subfolders = []
+    subfolders.append('None') 
+    path = directory
+    name_list = os.listdir(path)
+    full_list = [os.path.join(path,i) for i in name_list]
+    time_sorted_list = sorted(full_list, key=os.path.getmtime,reverse=True)
+
+    for entry in time_sorted_list:
+        if os.path.isdir(entry):
+            entry_str = f"{entry}"  # Convert entry to a string
+            full_path = entry_str
+            entry_str = entry_str.replace('\\','/')
+            entry_str = entry_str.replace(f"{directory}", "")  # Remove directory part
+            subfolders.append(entry_str)
+
+    return subfolders
+
+def get_available_loras_local(_sortedByTime):
+    
+    model_dir = shared.args.lora_dir  # Update with the appropriate directory path
+    subfolders = []
+    if _sortedByTime:
+        subfolders = list_subfoldersByTime(model_dir)
+    else:
+        subfolders = utils.get_available_loras()        
+
+    return subfolders
+
+
+# FPHAM SPLIT BY SENTENCE BLOCK ===============
+     
+def split_sentences(text: str, cutoff_len: int):
+    sentences = []
+    sentence = ''
+    delimiters = ['. ', '? ', '! ', '... ', '.\n', '?\n', '!\n','...\n','</s>','<//>']
+    abbreviations = ['Mr. ', 'Mrs. ', 'Dr. ', 'Ms. ', 'St. ', 'Prof. ', 'Jr. ', 'Ltd. ', 'Capt. ', 'Col. ', 'Gen. ', 'Ave. ', 'Blvd. ', 'Co. ', 'Corp. ', 'Dept. ', 'Est. ', 'Gov. ', 'Inc. ', 'Ph.D. ', 'Univ. ']
+    errors = 0
+    max_cut = cutoff_len-1
+    prev_char = ''  
+
+    for char in text:
+        sentence += char
+
+    
+        if (any(sentence.endswith(delimiter) for delimiter in delimiters) and
+            not (prev_char.isupper() and len(sentence) >= 3 and sentence[-3] != ' ') and 
+            not any(sentence.endswith(abbreviation) for abbreviation in abbreviations)):
+            tokens = shared.tokenizer.encode(sentence)
+            
+            if len(tokens) > max_cut:
+                tokens = tokens[:max_cut]
+                sentence = shared.tokenizer.decode(tokens, skip_special_tokens=True)
+                errors = errors + 1
+
+            sentences.append({'text': sentence, 'size': len(tokens)})
+            
+            sentence = ''
+
+        prev_char = char
+
+    if sentence:
+        tokens = shared.tokenizer.encode(sentence)
+        if len(tokens) > max_cut:
+            tokens = tokens[:max_cut]
+            sentence = shared.tokenizer.decode(tokens, skip_special_tokens=True)  
+            errors = errors + 1
+
+        sentences.append({'text': sentence, 'size': len(tokens)})
+
+    if errors > 0:
+        print(f"Trimmed sentences beyond Cutoff Length: {errors}")
+
+    return sentences
+
+# The goal of following code is to create blocks of text + overlapping blocks while:
+# respects sentence boundaries
+# always uses all the text 
+# hard cut defined by hard_cut_string or </s> will always end at the end of data block
+# no overlapping blocks will be created across hard cut or across </s> token
+
+def precise_cut(text: str, overlap: bool, min_chars_cut: int, eos_to_hc: bool, cutoff_len: int, hard_cut_string: str, debug_slicer:bool):
+
+    EOSX_str = '<//>' #hardcut placeholder
+    EOS_str = '</s>' 
+    print("Precise raw text slicer: ON")
+    
+    cut_string = hard_cut_string.replace('\\n', '\n')
+    text = text.replace(cut_string, EOSX_str)
+    sentences = split_sentences(text, cutoff_len)
+
+    print(f"Sentences: {len(sentences)}")
+    sentencelist = []
+    currentSentence = ''
+    totalLength = 0
+    max_cut = cutoff_len-1
+    half_cut = cutoff_len//2
+    halfcut_length = 0
+
+    edgeindex = []
+    half_index = 0
+
+    for index, item in enumerate(sentences):
+        
+        if halfcut_length+ item['size'] < half_cut:
+            halfcut_length += item['size']
+            half_index = index
+        else:
+            edgeindex.append(half_index)
+            halfcut_length = -2 * max_cut
+
+
+        if totalLength + item['size'] < max_cut and not currentSentence.endswith(EOSX_str): 
+            currentSentence += item['text']
+            totalLength += item['size']
+        else:
+
+            if len(currentSentence.strip()) > min_chars_cut:
+                sentencelist.append(currentSentence.strip())
+
+            currentSentence = item['text']
+            totalLength = item['size']
+            halfcut_length = item['size']
+            
+    if len(currentSentence.strip()) > min_chars_cut:    
+        sentencelist.append(currentSentence.strip())
+
+    unique_blocks = len(sentencelist)
+    print(f"Text Blocks: {unique_blocks}")
+
+    #overlap strategies: 
+    # don't overlap across HARD CUT (EOSX)
+    if overlap:
+        for edge_idx in edgeindex:
+            currentSentence = ''
+            totalLength = 0
+
+            for item in sentences[edge_idx:]:
+                if totalLength + item['size'] < max_cut:
+                    currentSentence += item['text']
+                    totalLength += item['size']
+                else:
+                    #if by chance EOSX is at the end then it's acceptable
+                    if currentSentence.endswith(EOSX_str) and len(currentSentence.strip()) > min_chars_cut:
+                            sentencelist.append(currentSentence.strip())    
+                    # otherwise don't cross hard cut    
+                    elif EOSX_str not in currentSentence and len(currentSentence.strip()) > min_chars_cut:
+                        sentencelist.append(currentSentence.strip())
+                    
+                    currentSentence = ''
+                    totalLength = 0
+                    break
+        
+        print(f"+ Overlapping blocks: {len(sentencelist)-unique_blocks}")
+
+    num_EOS = 0
+    for i in range(len(sentencelist)):
+        if eos_to_hc:
+            sentencelist[i] = sentencelist[i].replace(EOSX_str, EOS_str)
+        else:
+            sentencelist[i] = sentencelist[i].replace(EOSX_str, '')
+        
+        #someone may have had stop strings in the raw text...
+        sentencelist[i] = sentencelist[i].replace("</s></s>", EOS_str)
+        num_EOS += sentencelist[i].count(EOS_str)
+
+    if num_EOS > 0:
+        print(f"+ EOS count: {num_EOS}")
+
+    #final check for useless lines
+    sentencelist = [item for item in sentencelist if item.strip() != "</s>"]
+    sentencelist = [item for item in sentencelist if item.strip() != ""]
+
+
+    if debug_slicer:
+                    # Write the log file
+        Path('logs').mkdir(exist_ok=True)
+        sentencelist_dict = {index: sentence for index, sentence in enumerate(sentencelist)}
+        output_file = "logs/sentencelist.json"
+        with open(output_file, 'w') as f:
+            json.dump(sentencelist_dict, f,indent=2)
+        
+        print("Saved sentencelist.json in logs folder")
+    
+    return sentencelist   
+
+
+def sliding_block_cut(text: str, min_chars_cut: int, eos_to_hc: bool, cutoff_len: int, hard_cut_string: str, debug_slicer:bool):
+
+    EOSX_str = '<//>' #hardcut placeholder
+    EOS_str = '</s>' 
+    print("Mega Block Overlap: ON")
+    
+    cut_string = hard_cut_string.replace('\\n', '\n')
+    text = text.replace(cut_string, EOSX_str)
+    sentences = split_sentences(text, cutoff_len)
+
+    print(f"Sentences: {len(sentences)}")
+    sentencelist = []
+    
+    max_cut = cutoff_len-1
+
+    #print(f"max_cut: {max_cut}")
+    advancing_to = 0
+
+    prev_block_lastsentence = ""
+    
+
+    for i in range(len(sentences)):
+        totalLength = 0
+        currentSentence = ''
+        lastsentence = ""
+        
+        if i >= advancing_to:
+            for k in range(i, len(sentences)):
+                
+                current_length = sentences[k]['size']
+
+                if totalLength + current_length <= max_cut and not currentSentence.endswith(EOSX_str):
+                    currentSentence += sentences[k]['text']
+                    totalLength += current_length
+                    lastsentence = sentences[k]['text']
+                else:
+                    if len(currentSentence.strip()) > min_chars_cut:
+                        if prev_block_lastsentence!=lastsentence:
+                            sentencelist.append(currentSentence.strip())
+                            prev_block_lastsentence = lastsentence
+                        
+                    advancing_to = 0
+                    if currentSentence.endswith(EOSX_str):
+                        advancing_to = k
+
+                    currentSentence = ""
+                    totalLength = 0
+                    break
+            
+            if currentSentence != "":
+                if len(currentSentence.strip()) > min_chars_cut:
+                    sentencelist.append(currentSentence.strip())
+
+    unique_blocks = len(sentencelist)
+    print(f"Text Blocks: {unique_blocks}")
+    num_EOS = 0
+    for i in range(len(sentencelist)):
+        if eos_to_hc:
+            sentencelist[i] = sentencelist[i].replace(EOSX_str, EOS_str)
+        else:
+            sentencelist[i] = sentencelist[i].replace(EOSX_str, '')
+        
+        #someone may have had stop strings in the raw text...
+        sentencelist[i] = sentencelist[i].replace("</s></s>", EOS_str)
+        num_EOS += sentencelist[i].count(EOS_str)
+
+    if num_EOS > 0:
+        print(f"+ EOS count: {num_EOS}")
+
+    #final check for useless lines
+    sentencelist = [item for item in sentencelist if item.strip() != "</s>"]
+    sentencelist = [item for item in sentencelist if item.strip() != ""]
+
+
+    if debug_slicer:
+                    # Write the log file
+        Path('logs').mkdir(exist_ok=True)
+        sentencelist_dict = {index: sentence for index, sentence in enumerate(sentencelist)}
+        output_file = "logs/sentencelist.json"
+        with open(output_file, 'w') as f:
+            json.dump(sentencelist_dict, f,indent=2)
+        
+        print("Saved sentencelist.json in logs folder")
+    
+    return sentencelist   
+
+# Example usage:
+# download_file_from_url('https://example.com/path/to/your/file.ext', '/output/directory')
+
+def download_file_from_url(url, overwrite, output_dir_in, valid_extensions = {'.txt', '.json'}):
+    try:
+    # Validate and sanitize the URL
+    #parsed_url = urllib.parse.urlparse(url)
+    #if not parsed_url.netloc:
+    #    raise ValueError("Invalid URL")
+    #filename = os.path.basename(parsed_url.path)
+
+    # Get the filename from the URL
+
+        session = requests.Session()
+        headers = {}
+        mode = 'wb'
+        filename = url.split('/')[-1]
+
+        output_dir = str(output_dir_in)
+        # Construct the full path to the output file
+        local_filename = os.path.join(output_dir, filename)
+
+        # Check if the local file already exists
+        overw = ''
+        if os.path.exists(local_filename):
+            if not overwrite:
+                yield f"File '{local_filename}' already exists. Aborting."
+                return
+            else:
+                overw = ' [Overwrite existing]'
+
+        filename_lower = filename.lower()
+
+        # Send an HTTP GET request to the URL with a timeout
+        file_extension = os.path.splitext(filename_lower)[-1]
+        
+        if file_extension not in valid_extensions:
+            yield f"Invalid file extension: {file_extension}. Only {valid_extensions} files are supported."
+            return
+
+        with session.get(url, stream=True, headers=headers, timeout=10) as r:
+            r.raise_for_status() 
+            # total size can be wildly inaccurate
+            #total_size = int(r.headers.get('content-length', 0))
+            
+            block_size = 1024 * 4  
+            with open(local_filename, mode) as f:
+                count = 0
+                for data in r.iter_content(block_size):
+                    f.write(data)
+                    count += len(data)
+
+                    yield f"Downloaded: {count} " + overw
+
+            # Verify file size if possible
+            if os.path.exists(local_filename):
+                downloaded_size = os.path.getsize(local_filename)
+                if downloaded_size > 0:
+                    yield f"File '{filename}' downloaded to '{output_dir}' ({downloaded_size} bytes)."
+                    print("File Downloaded")
+                else:
+                    print("Downloaded file is zero")
+                    yield f"Failed. Downloaded file size is zero)."
+            else:
+                print(f"Error: {local_filename} failed to download.")
+                yield f"Error: {local_filename} failed to download"
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        yield f"An error occurred: {e}"
+
+    finally:
+        # Close the session to release resources
+        session.close()
+

extensions/character_bias/script.py

@@ -0,0 +1,74 @@
+import os
+
+import gradio as gr
+
+# get the current directory of the script
+current_dir = os.path.dirname(os.path.abspath(__file__))
+
+# check if the bias_options.txt file exists, if not, create it
+bias_file = os.path.join(current_dir, "bias_options.txt")
+if not os.path.isfile(bias_file):
+    with open(bias_file, "w") as f:
+        f.write("*I am so happy*\n*I am so sad*\n*I am so excited*\n*I am so bored*\n*I am so angry*")
+
+# read bias options from the text file
+with open(bias_file, "r") as f:
+    bias_options = [line.strip() for line in f.readlines()]
+
+params = {
+    "activate": True,
+    "bias string": " *I am so happy*",
+    "custom string": "",
+}
+
+
+def input_modifier(string):
+    """
+    This function is applied to your text inputs before
+    they are fed into the model.
+    """
+    return string
+
+
+def output_modifier(string):
+    """
+    This function is applied to the model outputs.
+    """
+    return string
+
+
+def bot_prefix_modifier(string):
+    """
+    This function is only applied in chat mode. It modifies
+    the prefix text for the Bot and can be used to bias its
+    behavior.
+    """
+    if params['activate']:
+        if params['custom string'].strip() != '':
+            return f'{string} {params["custom string"].strip()} '
+        else:
+            return f'{string} {params["bias string"].strip()} '
+    else:
+        return string
+
+
+def ui():
+    # Gradio elements
+    activate = gr.Checkbox(value=params['activate'], label='Activate character bias')
+    dropdown_string = gr.Dropdown(choices=bias_options, value=params["bias string"], label='Character bias', info='To edit the options in this dropdown edit the "bias_options.txt" file')
+    custom_string = gr.Textbox(value=params['custom string'], placeholder="Enter custom bias string", label="Custom Character Bias", info='If not empty, will be used instead of the value above')
+
+    # Event functions to update the parameters in the backend
+    def update_bias_string(x):
+        if x:
+            params.update({"bias string": x})
+        else:
+            params.update({"bias string": dropdown_string.get()})
+        return x
+
+    def update_custom_string(x):
+        params.update({"custom string": x})
+
+    dropdown_string.change(update_bias_string, dropdown_string, None)
+    custom_string.change(update_custom_string, custom_string, None)
+    activate.change(lambda x: params.update({"activate": x}), activate, None)

extensions/coqui_tts/harvard_sentences.txt

@@ -0,0 +1,720 @@
+The birch canoe slid on the smooth planks.
+Glue the sheet to the dark blue background.
+It's easy to tell the depth of a well.
+These days a chicken leg is a rare dish.
+Rice is often served in round bowls.
+The juice of lemons makes fine punch.
+The box was thrown beside the parked truck.
+The hogs were fed chopped corn and garbage.
+Four hours of steady work faced us.
+A large size in stockings is hard to sell.
+The boy was there when the sun rose.
+A rod is used to catch pink salmon.
+The source of the huge river is the clear spring.
+Kick the ball straight and follow through.
+Help the woman get back to her feet.
+A pot of tea helps to pass the evening.
+Smoky fires lack flame and heat.
+The soft cushion broke the man's fall.
+The salt breeze came across from the sea.
+The girl at the booth sold fifty bonds.
+The small pup gnawed a hole in the sock.
+The fish twisted and turned on the bent hook.
+Press the pants and sew a button on the vest.
+The swan dive was far short of perfect.
+The beauty of the view stunned the young boy.
+Two blue fish swam in the tank.
+Her purse was full of useless trash.
+The colt reared and threw the tall rider.
+It snowed, rained, and hailed the same morning.
+Read verse out loud for pleasure.
+Hoist the load to your left shoulder.
+Take the winding path to reach the lake.
+Note closely the size of the gas tank.
+Wipe the grease off his dirty face.
+Mend the coat before you go out.
+The wrist was badly strained and hung limp.
+The stray cat gave birth to kittens.
+The young girl gave no clear response.
+The meal was cooked before the bell rang.
+What joy there is in living.
+A king ruled the state in the early days.
+The ship was torn apart on the sharp reef.
+Sickness kept him home the third week.
+The wide road shimmered in the hot sun.
+The lazy cow lay in the cool grass.
+Lift the square stone over the fence.
+The rope will bind the seven books at once.
+Hop over the fence and plunge in.
+The friendly gang left the drug store.
+Mesh wire keeps chicks inside.
+The frosty air passed through the coat.
+The crooked maze failed to fool the mouse.
+Adding fast leads to wrong sums.
+The show was a flop from the very start.
+A saw is a tool used for making boards.
+The wagon moved on well oiled wheels.
+March the soldiers past the next hill.
+A cup of sugar makes sweet fudge.
+Place a rosebush near the porch steps.
+Both lost their lives in the raging storm.
+We talked of the side show in the circus.
+Use a pencil to write the first draft.
+He ran half way to the hardware store.
+The clock struck to mark the third period.
+A small creek cut across the field.
+Cars and busses stalled in snow drifts.
+The set of china hit the floor with a crash.
+This is a grand season for hikes on the road.
+The dune rose from the edge of the water.
+Those words were the cue for the actor to leave.
+A yacht slid around the point into the bay.
+The two met while playing on the sand.
+The ink stain dried on the finished page.
+The walled town was seized without a fight.
+The lease ran out in sixteen weeks.
+A tame squirrel makes a nice pet.
+The horn of the car woke the sleeping cop.
+The heart beat strongly and with firm strokes.
+The pearl was worn in a thin silver ring.
+The fruit peel was cut in thick slices.
+The Navy attacked the big task force.
+See the cat glaring at the scared mouse.
+There are more than two factors here.
+The hat brim was wide and too droopy.
+The lawyer tried to lose his case.
+The grass curled around the fence post.
+Cut the pie into large parts.
+Men strive but seldom get rich.
+Always close the barn door tight.
+He lay prone and hardly moved a limb.
+The slush lay deep along the street.
+A wisp of cloud hung in the blue air.
+A pound of sugar costs more than eggs.
+The fin was sharp and cut the clear water.
+The play seems dull and quite stupid.
+Bail the boat to stop it from sinking.
+The term ended in late June that year.
+A tusk is used to make costly gifts.
+Ten pins were set in order.
+The bill was paid every third week.
+Oak is strong and also gives shade.
+Cats and dogs each hate the other.
+The pipe began to rust while new.
+Open the crate but don't break the glass.
+Add the sum to the product of these three.
+Thieves who rob friends deserve jail.
+The ripe taste of cheese improves with age.
+Act on these orders with great speed.
+The hog crawled under the high fence.
+Move the vat over the hot fire.
+The bark of the pine tree was shiny and dark.
+Leaves turn brown and yellow in the fall.
+The pennant waved when the wind blew.
+Split the log with a quick, sharp blow.
+Burn peat after the logs give out.
+He ordered peach pie with ice cream.
+Weave the carpet on the right hand side.
+Hemp is a weed found in parts of the tropics.
+A lame back kept his score low.
+We find joy in the simplest things.
+Type out three lists of orders.
+The harder he tried the less he got done.
+The boss ran the show with a watchful eye.
+The cup cracked and spilled its contents.
+Paste can cleanse the most dirty brass.
+The slang word for raw whiskey is booze.
+It caught its hind paw in a rusty trap.
+The wharf could be seen at the farther shore.
+Feel the heat of the weak dying flame.
+The tiny girl took off her hat.
+A cramp is no small danger on a swim.
+He said the same phrase thirty times.
+Pluck the bright rose without leaves.
+Two plus seven is less than ten.
+The glow deepened in the eyes of the sweet girl.
+Bring your problems to the wise chief.
+Write a fond note to the friend you cherish.
+Clothes and lodging are free to new men.
+We frown when events take a bad turn.
+Port is a strong wine with a smoky taste.
+The young kid jumped the rusty gate.
+Guess the results from the first scores.
+A salt pickle tastes fine with ham.
+The just claim got the right verdict.
+These thistles bend in a high wind.
+Pure bred poodles have curls.
+The tree top waved in a graceful way.
+The spot on the blotter was made by green ink.
+Mud was spattered on the front of his white shirt.
+The cigar burned a hole in the desk top.
+The empty flask stood on the tin tray.
+A speedy man can beat this track mark.
+He broke a new shoelace that day.
+The coffee stand is too high for the couch.
+The urge to write short stories is rare.
+The pencils have all been used.
+The pirates seized the crew of the lost ship.
+We tried to replace the coin but failed.
+She sewed the torn coat quite neatly.
+The sofa cushion is red and of light weight.
+The jacket hung on the back of the wide chair.
+At that high level the air is pure.
+Drop the two when you add the figures.
+A filing case is now hard to buy.
+An abrupt start does not win the prize.
+Wood is best for making toys and blocks.
+The office paint was a dull, sad tan.
+He knew the skill of the great young actress.
+A rag will soak up spilled water.
+A shower of dirt fell from the hot pipes.
+Steam hissed from the broken valve.
+The child almost hurt the small dog.
+There was a sound of dry leaves outside.
+The sky that morning was clear and bright blue.
+Torn scraps littered the stone floor.
+Sunday is the best part of the week.
+The doctor cured him with these pills.
+The new girl was fired today at noon.
+They felt gay when the ship arrived in port.
+Add the store's account to the last cent.
+Acid burns holes in wool cloth.
+Fairy tales should be fun to write.
+Eight miles of woodland burned to waste.
+The third act was dull and tired the players.
+A young child should not suffer fright.
+Add the column and put the sum here.
+We admire and love a good cook.
+There the flood mark is ten inches.
+He carved a head from the round block of marble.
+She has a smart way of wearing clothes.
+The fruit of a fig tree is apple-shaped.
+Corn cobs can be used to kindle a fire.
+Where were they when the noise started.
+The paper box is full of thumb tacks.
+Sell your gift to a buyer at a good gain.
+The tongs lay beside the ice pail.
+The petals fall with the next puff of wind.
+Bring your best compass to the third class.
+They could laugh although they were sad.
+Farmers came in to thresh the oat crop.
+The brown house was on fire to the attic.
+The lure is used to catch trout and flounder.
+Float the soap on top of the bath water.
+A blue crane is a tall wading bird.
+A fresh start will work such wonders.
+The club rented the rink for the fifth night.
+After the dance, they went straight home.
+The hostess taught the new maid to serve.
+He wrote his last novel there at the inn.
+Even the worst will beat his low score.
+The cement had dried when he moved it.
+The loss of the second ship was hard to take.
+The fly made its way along the wall.
+Do that with a wooden stick.
+Live wires should be kept covered.
+The large house had hot water taps.
+It is hard to erase blue or red ink.
+Write at once or you may forget it.
+The doorknob was made of bright clean brass.
+The wreck occurred by the bank on Main Street.
+A pencil with black lead writes best.
+Coax a young calf to drink from a bucket.
+Schools for ladies teach charm and grace.
+The lamp shone with a steady green flame.
+They took the axe and the saw to the forest.
+The ancient coin was quite dull and worn.
+The shaky barn fell with a loud crash.
+Jazz and swing fans like fast music.
+Rake the rubbish up and then burn it.
+Slash the gold cloth into fine ribbons.
+Try to have the court decide the case.
+They are pushed back each time they attack.
+He broke his ties with groups of former friends.
+They floated on the raft to sun their white backs.
+The map had an X that meant nothing.
+Whitings are small fish caught in nets.
+Some ads serve to cheat buyers.
+Jerk the rope and the bell rings weakly.
+A waxed floor makes us lose balance.
+Madam, this is the best brand of corn.
+On the islands the sea breeze is soft and mild.
+The play began as soon as we sat down.
+This will lead the world to more sound and fury.
+Add salt before you fry the egg.
+The rush for funds reached its peak Tuesday.
+The birch looked stark white and lonesome.
+The box is held by a bright red snapper.
+To make pure ice, you freeze water.
+The first worm gets snapped early.
+Jump the fence and hurry up the bank.
+Yell and clap as the curtain slides back.
+They are men who walk the middle of the road.
+Both brothers wear the same size.
+In some form or other we need fun.
+The prince ordered his head chopped off.
+The houses are built of red clay bricks.
+Ducks fly north but lack a compass.
+Fruit flavors are used in fizz drinks.
+These pills do less good than others.
+Canned pears lack full flavor.
+The dark pot hung in the front closet.
+Carry the pail to the wall and spill it there.
+The train brought our hero to the big town.
+We are sure that one war is enough.
+Gray paint stretched for miles around.
+The rude laugh filled the empty room.
+High seats are best for football fans.
+Tea served from the brown jug is tasty.
+A dash of pepper spoils beef stew.
+A zestful food is the hot-cross bun.
+The horse trotted around the field at a brisk pace.
+Find the twin who stole the pearl necklace.
+Cut the cord that binds the box tightly.
+The red tape bound the smuggled food.
+Look in the corner to find the tan shirt.
+The cold drizzle will halt the bond drive.
+Nine men were hired to dig the ruins.
+The junk yard had a mouldy smell.
+The flint sputtered and lit a pine torch.
+Soak the cloth and drown the sharp odor.
+The shelves were bare of both jam or crackers.
+A joy to every child is the swan boat.
+All sat frozen and watched the screen.
+A cloud of dust stung his tender eyes.
+To reach the end he needs much courage.
+Shape the clay gently into block form.
+A ridge on a smooth surface is a bump or flaw.
+Hedge apples may stain your hands green.
+Quench your thirst, then eat the crackers.
+Tight curls get limp on rainy days.
+The mute muffled the high tones of the horn.
+The gold ring fits only a pierced ear.
+The old pan was covered with hard fudge.
+Watch the log float in the wide river.
+The node on the stalk of wheat grew daily.
+The heap of fallen leaves was set on fire.
+Write fast if you want to finish early.
+His shirt was clean but one button was gone.
+The barrel of beer was a brew of malt and hops.
+Tin cans are absent from store shelves.
+Slide the box into that empty space.
+The plant grew large and green in the window.
+The beam dropped down on the workmen's head.
+Pink clouds floated with the breeze.
+She danced like a swan, tall and graceful.
+The tube was blown and the tire flat and useless.
+It is late morning on the old wall clock.
+Let's all join as we sing the last chorus.
+The last switch cannot be turned off.
+The fight will end in just six minutes.
+The store walls were lined with colored frocks.
+The peace league met to discuss their plans.
+The rise to fame of a person takes luck.
+Paper is scarce, so write with much care.
+The quick fox jumped on the sleeping cat.
+The nozzle of the fire hose was bright brass.
+Screw the round cap on as tight as needed.
+Time brings us many changes.
+The purple tie was ten years old.
+Men think and plan and sometimes act.
+Fill the ink jar with sticky glue.
+He smoke a big pipe with strong contents.
+We need grain to keep our mules healthy.
+Pack the records in a neat thin case.
+The crunch of feet in the snow was the only sound.
+The copper bowl shone in the sun's rays.
+Boards will warp unless kept dry.
+The plush chair leaned against the wall.
+Glass will clink when struck by metal.
+Bathe and relax in the cool green grass.
+Nine rows of soldiers stood in line.
+The beach is dry and shallow at low tide.
+The idea is to sew both edges straight.
+The kitten chased the dog down the street.
+Pages bound in cloth make a book.
+Try to trace the fine lines of the painting.
+Women form less than half of the group.
+The zones merge in the central part of town.
+A gem in the rough needs work to polish.
+Code is used when secrets are sent.
+Most of the news is easy for us to hear.
+He used the lathe to make brass objects.
+The vane on top of the pole revolved in the wind.
+Mince pie is a dish served to children.
+The clan gathered on each dull night.
+Let it burn, it gives us warmth and comfort.
+A castle built from sand fails to endure.
+A child's wit saved the day for us.
+Tack the strip of carpet to the worn floor.
+Next Tuesday we must vote.
+Pour the stew from the pot into the plate.
+Each penny shone like new.
+The man went to the woods to gather sticks.
+The dirt piles were lines along the road.
+The logs fell and tumbled into the clear stream.
+Just hoist it up and take it away.
+A ripe plum is fit for a king's palate.
+Our plans right now are hazy.
+Brass rings are sold by these natives.
+It takes a good trap to capture a bear.
+Feed the white mouse some flower seeds.
+The thaw came early and freed the stream.
+He took the lead and kept it the whole distance.
+The key you designed will fit the lock.
+Plead to the council to free the poor thief.
+Better hash is made of rare beef.
+This plank was made for walking on.
+The lake sparkled in the red hot sun.
+He crawled with care along the ledge.
+Tend the sheep while the dog wanders.
+It takes a lot of help to finish these.
+Mark the spot with a sign painted red.
+Take two shares as a fair profit.
+The fur of cats goes by many names.
+North winds bring colds and fevers.
+He asks no person to vouch for him.
+Go now and come here later.
+A sash of gold silk will trim her dress.
+Soap can wash most dirt away.
+That move means the game is over.
+He wrote down a long list of items.
+A siege will crack the strong defense.
+Grape juice and water mix well.
+Roads are paved with sticky tar.
+Fake stones shine but cost little.
+The drip of the rain made a pleasant sound.
+Smoke poured out of every crack.
+Serve the hot rum to the tired heroes.
+Much of the story makes good sense.
+The sun came up to light the eastern sky.
+Heave the line over the port side.
+A lathe cuts and trims any wood.
+It's a dense crowd in two distinct ways.
+His hip struck the knee of the next player.
+The stale smell of old beer lingers.
+The desk was firm on the shaky floor.
+It takes heat to bring out the odor.
+Beef is scarcer than some lamb.
+Raise the sail and steer the ship northward.
+A cone costs five cents on Mondays.
+A pod is what peas always grow in.
+Jerk the dart from the cork target.
+No cement will hold hard wood.
+We now have a new base for shipping.
+A list of names is carved around the base.
+The sheep were led home by a dog.
+Three for a dime, the young peddler cried.
+The sense of smell is better than that of touch.
+No hardship seemed to keep him sad.
+Grace makes up for lack of beauty.
+Nudge gently but wake her now.
+The news struck doubt into restless minds.
+Once we stood beside the shore.
+A chink in the wall allowed a draft to blow.
+Fasten two pins on each side.
+A cold dip restores health and zest.
+He takes the oath of office each March.
+The sand drifts over the sill of the old house.
+The point of the steel pen was bent and twisted.
+There is a lag between thought and act.
+Seed is needed to plant the spring corn.
+Draw the chart with heavy black lines.
+The boy owed his pal thirty cents.
+The chap slipped into the crowd and was lost.
+Hats are worn to tea and not to dinner.
+The ramp led up to the wide highway.
+Beat the dust from the rug onto the lawn.
+Say it slowly but make it ring clear.
+The straw nest housed five robins.
+Screen the porch with woven straw mats.
+This horse will nose his way to the finish.
+The dry wax protects the deep scratch.
+He picked up the dice for a second roll.
+These coins will be needed to pay his debt.
+The nag pulled the frail cart along.
+Twist the valve and release hot steam.
+The vamp of the shoe had a gold buckle.
+The smell of burned rags itches my nose.
+New pants lack cuffs and pockets.
+The marsh will freeze when cold enough.
+They slice the sausage thin with a knife.
+The bloom of the rose lasts a few days.
+A gray mare walked before the colt.
+Breakfast buns are fine with a hot drink.
+Bottles hold four kinds of rum.
+The man wore a feather in his felt hat.
+He wheeled the bike past the winding road.
+Drop the ashes on the worn old rug.
+The desk and both chairs were painted tan.
+Throw out the used paper cup and plate.
+A clean neck means a neat collar.
+The couch cover and hall drapes were blue.
+The stems of the tall glasses cracked and broke.
+The wall phone rang loud and often.
+The clothes dried on a thin wooden rack.
+Turn on the lantern which gives us light.
+The cleat sank deeply into the soft turf.
+The bills were mailed promptly on the tenth of the month.
+To have is better than to wait and hope.
+The price is fair for a good antique clock.
+The music played on while they talked.
+Dispense with a vest on a day like this.
+The bunch of grapes was pressed into wine.
+He sent the figs, but kept the ripe cherries.
+The hinge on the door creaked with old age.
+The screen before the fire kept in the sparks.
+Fly by night, and you waste little time.
+Thick glasses helped him read the print.
+Birth and death mark the limits of life.
+The chair looked strong but had no bottom.
+The kite flew wildly in the high wind.
+A fur muff is stylish once more.
+The tin box held priceless stones.
+We need an end of all such matter.
+The case was puzzling to the old and wise.
+The bright lanterns were gay on the dark lawn.
+We don't get much money but we have fun.
+The youth drove with zest, but little skill.
+Five years he lived with a shaggy dog.
+A fence cuts through the corner lot.
+The way to save money is not to spend much.
+Shut the hatch before the waves push it in.
+The odor of spring makes young hearts jump.
+Crack the walnut with your sharp side teeth.
+He offered proof in the form of a large chart.
+Send the stuff in a thick paper bag.
+A quart of milk is water for the most part.
+They told wild tales to frighten him.
+The three story house was built of stone.
+In the rear of the ground floor was a large passage.
+A man in a blue sweater sat at the desk.
+Oats are a food eaten by horse and man.
+Their eyelids droop for want of sleep.
+A sip of tea revives his tired friend.
+There are many ways to do these things.
+Tuck the sheet under the edge of the mat.
+A force equal to that would move the earth.
+We like to see clear weather.
+The work of the tailor is seen on each side.
+Take a chance and win a china doll.
+Shake the dust from your shoes, stranger.
+She was kind to sick old people.
+The square wooden crate was packed to be shipped.
+The dusty bench stood by the stone wall.
+We dress to suit the weather of most days.
+Smile when you say nasty words.
+A bowl of rice is free with chicken stew.
+The water in this well is a source of good health.
+Take shelter in this tent, but keep still.
+That guy is the writer of a few banned books.
+The little tales they tell are false.
+The door was barred, locked, and bolted as well.
+Ripe pears are fit for a queen's table.
+A big wet stain was on the round carpet.
+The kite dipped and swayed, but stayed aloft.
+The pleasant hours fly by much too soon.
+The room was crowded with a wild mob.
+This strong arm shall shield your honor.
+She blushed when he gave her a white orchid.
+The beetle droned in the hot June sun.
+Press the pedal with your left foot.
+Neat plans fail without luck.
+The black trunk fell from the landing.
+The bank pressed for payment of the debt.
+The theft of the pearl pin was kept secret.
+Shake hands with this friendly child.
+The vast space stretched into the far distance.
+A rich farm is rare in this sandy waste.
+His wide grin earned many friends.
+Flax makes a fine brand of paper.
+Hurdle the pit with the aid of a long pole.
+A strong bid may scare your partner stiff.
+Even a just cause needs power to win.
+Peep under the tent and see the clowns.
+The leaf drifts along with a slow spin.
+Cheap clothes are flashy but don't last.
+A thing of small note can cause despair.
+Flood the mails with requests for this book.
+A thick coat of black paint covered all.
+The pencil was cut to be sharp at both ends.
+Those last words were a strong statement.
+He wrote his name boldly at the top of the sheet.
+Dill pickles are sour but taste fine.
+Down that road is the way to the grain farmer.
+Either mud or dust are found at all times.
+The best method is to fix it in place with clips.
+If you mumble your speech will be lost.
+At night the alarm roused him from a deep sleep.
+Read just what the meter says.
+Fill your pack with bright trinkets for the poor.
+The small red neon lamp went out.
+Clams are small, round, soft, and tasty.
+The fan whirled its round blades softly.
+The line where the edges join was clean.
+Breathe deep and smell the piny air.
+It matters not if he reads these words or those.
+A brown leather bag hung from its strap.
+A toad and a frog are hard to tell apart.
+A white silk jacket goes with any shoes.
+A break in the dam almost caused a flood.
+Paint the sockets in the wall dull green.
+The child crawled into the dense grass.
+Bribes fail where honest men work.
+Trample the spark, else the flames will spread.
+The hilt of the sword was carved with fine designs.
+A round hole was drilled through the thin board.
+Footprints showed the path he took up the beach.
+She was waiting at my front lawn.
+A vent near the edge brought in fresh air.
+Prod the old mule with a crooked stick.
+It is a band of steel three inches wide.
+The pipe ran almost the length of the ditch.
+It was hidden from sight by a mass of leaves and shrubs.
+The weight of the package was seen on the high scale.
+Wake and rise, and step into the green outdoors.
+The green light in the brown box flickered.
+The brass tube circled the high wall.
+The lobes of her ears were pierced to hold rings.
+Hold the hammer near the end to drive the nail.
+Next Sunday is the twelfth of the month.
+Every word and phrase he speaks is true.
+He put his last cartridge into the gun and fired.
+They took their kids from the public school.
+Drive the screw straight into the wood.
+Keep the hatch tight and the watch constant.
+Sever the twine with a quick snip of the knife.
+Paper will dry out when wet.
+Slide the catch back and open the desk.
+Help the weak to preserve their strength.
+A sullen smile gets few friends.
+Stop whistling and watch the boys march.
+Jerk the cord, and out tumbles the gold.
+Slide the tray across the glass top.
+The cloud moved in a stately way and was gone.
+Light maple makes for a swell room.
+Set the piece here and say nothing.
+Dull stories make her laugh.
+A stiff cord will do to fasten your shoe.
+Get the trust fund to the bank early.
+Choose between the high road and the low.
+A plea for funds seems to come again.
+He lent his coat to the tall gaunt stranger.
+There is a strong chance it will happen once more.
+The duke left the park in a silver coach.
+Greet the new guests and leave quickly.
+When the frost has come it is time for turkey.
+Sweet words work better than fierce.
+A thin stripe runs down the middle.
+A six comes up more often than a ten.
+Lush fern grow on the lofty rocks.
+The ram scared the school children off.
+The team with the best timing looks good.
+The farmer swapped his horse for a brown ox.
+Sit on the perch and tell the others what to do.
+A steep trail is painful for our feet.
+The early phase of life moves fast.
+Green moss grows on the northern side.
+Tea in thin china has a sweet taste.
+Pitch the straw through the door of the stable.
+The latch on the back gate needed a nail.
+The goose was brought straight from the old market.
+The sink is the thing in which we pile dishes.
+A whiff of it will cure the most stubborn cold.
+The facts don't always show who is right.
+She flaps her cape as she parades the street.
+The loss of the cruiser was a blow to the fleet.
+Loop the braid to the left and then over.
+Plead with the lawyer to drop the lost cause.
+Calves thrive on tender spring grass.
+Post no bills on this office wall.
+Tear a thin sheet from the yellow pad.
+A cruise in warm waters in a sleek yacht is fun.
+A streak of color ran down the left edge.
+It was done before the boy could see it.
+Crouch before you jump or miss the mark.
+Pack the kits and don't forget the salt.
+The square peg will settle in the round hole.
+Fine soap saves tender skin.
+Poached eggs and tea must suffice.
+Bad nerves are jangled by a door slam.
+Ship maps are different from those for planes.
+Dimes showered down from all sides.
+They sang the same tunes at each party.
+The sky in the west is tinged with orange red.
+The pods of peas ferment in bare fields.
+The horse balked and threw the tall rider.
+The hitch between the horse and cart broke.
+Pile the coal high in the shed corner.
+A gold vase is both rare and costly.
+The knife was hung inside its bright sheath.
+The rarest spice comes from the far East.
+The roof should be tilted at a sharp slant.
+A smatter of French is worse than none.
+The mule trod the treadmill day and night.
+The aim of the contest is to raise a great fund.
+To send it now in large amounts is bad.
+There is a fine hard tang in salty air.
+Cod is the main business of the north shore.
+The slab was hewn from heavy blocks of slate.
+Dunk the stale biscuits into strong drink.
+Hang tinsel from both branches.
+Cap the jar with a tight brass cover.
+The poor boy missed the boat again.
+Be sure to set the lamp firmly in the hole.
+Pick a card and slip it under the pack.
+A round mat will cover the dull spot.
+The first part of the plan needs changing.
+A good book informs of what we ought to know.
+The mail comes in three batches per day.
+You cannot brew tea in a cold pot.
+Dots of light betrayed the black cat.
+Put the chart on the mantel and tack it down.
+The night shift men rate extra pay.
+The red paper brightened the dim stage.
+See the player scoot to third base.
+Slide the bill between the two leaves.
+Many hands help get the job done.
+We don't like to admit our small faults.
+No doubt about the way the wind blows.
+Dig deep in the earth for pirate's gold.
+The steady drip is worse than a drenching rain.
+A flat pack takes less luggage space.
+Green ice frosted the punch bowl.
+A stuffed chair slipped from the moving van.
+The stitch will serve but needs to be shortened.
+A thin book fits in the side pocket.
+The gloss on top made it unfit to read.
+The hail pattered on the burnt brown grass.
+Seven seals were stamped on great sheets.
+Our troops are set to strike heavy blows.
+The store was jammed before the sale could start.
+It was a bad error on the part of the new judge.
+One step more and the board will collapse.
+Take the match and strike it against your shoe.
+The pot boiled, but the contents failed to jell.
+The baby puts his right foot in his mouth.
+The bombs left most of the town in ruins.
+Stop and stare at the hard working man.
+The streets are narrow and full of sharp turns.
+The pup jerked the leash as he saw a feline shape.
+Open your book to the first page.
+Fish evade the net and swim off.
+Dip the pail once and let it settle.
+Will you please answer that phone.
+The big red apple fell to the ground.
+The curtain rose and the show was on.
+The young prince became heir to the throne.
+He sent the boy on a short errand.
+Leave now and you will arrive on time.
+The corner store was robbed last night.
+A gold ring will please most any girl.
+The long journey home took a year.
+She saw a cat in the neighbor's house.
+A pink shell was found on the sandy beach.
+Small children came to see him.
+The grass and bushes were wet with dew.
+The blind man counted his old coins.
+A severe storm tore down the barn.
+She called his name many times.
+When you hear the bell, come quickly.

extensions/coqui_tts/languages.json

@@ -0,0 +1,18 @@
+{
+  "Arabic": "ar",
+  "Chinese": "zh-cn",
+  "Czech": "cs",
+  "Dutch": "nl",
+  "English": "en",
+  "French": "fr",
+  "German": "de",
+  "Hungarian": "hu",
+  "Italian": "it",
+  "Japanese": "ja",
+  "Korean": "ko",
+  "Polish": "pl",
+  "Portuguese": "pt",
+  "Russian": "ru",
+  "Spanish": "es",
+  "Turkish": "tr"
+}

extensions/coqui_tts/requirements.txt

@@ -0,0 +1 @@
+TTS==0.21.*

extensions/coqui_tts/script.py

@@ -0,0 +1,227 @@
+import html
+import json
+import os
+import random
+import time
+from pathlib import Path
+
+import gradio as gr
+from TTS.api import TTS
+from TTS.utils.synthesizer import Synthesizer
+
+from modules import chat, shared, ui_chat
+from modules.ui import create_refresh_button
+from modules.utils import gradio
+
+os.environ["COQUI_TOS_AGREED"] = "1"
+
+params = {
+    "activate": True,
+    "autoplay": True,
+    "show_text": False,
+    "remove_trailing_dots": False,
+    "voice": "female_01.wav",
+    "language": "English",
+    "model_name": "tts_models/multilingual/multi-dataset/xtts_v2",
+    "device": "cuda"
+}
+
+this_dir = str(Path(__file__).parent.resolve())
+model = None
+with open(Path(f"{this_dir}/languages.json"), encoding='utf8') as f:
+    languages = json.load(f)
+
+
+def get_available_voices():
+    return sorted([voice.name for voice in Path(f"{this_dir}/voices").glob("*.wav")])
+
+
+def preprocess(raw_input):
+    raw_input = html.unescape(raw_input)
+    # raw_input = raw_input.strip("\"")
+    return raw_input
+
+
+def new_split_into_sentences(self, text):
+    sentences = self.seg.segment(text)
+    if params['remove_trailing_dots']:
+        sentences_without_dots = []
+        for sentence in sentences:
+            if sentence.endswith('.') and not sentence.endswith('...'):
+                sentence = sentence[:-1]
+
+            sentences_without_dots.append(sentence)
+
+        return sentences_without_dots
+    else:
+        return sentences
+
+
+Synthesizer.split_into_sentences = new_split_into_sentences
+
+
+def load_model():
+    model = TTS(params["model_name"]).to(params["device"])
+    return model
+
+
+def remove_tts_from_history(history):
+    for i, entry in enumerate(history['internal']):
+        history['visible'][i] = [history['visible'][i][0], entry[1]]
+
+    return history
+
+
+def toggle_text_in_history(history):
+    for i, entry in enumerate(history['visible']):
+        visible_reply = entry[1]
+        if visible_reply.startswith('<audio'):
+            if params['show_text']:
+                reply = history['internal'][i][1]
+                history['visible'][i] = [history['visible'][i][0], f"{visible_reply.split('</audio>')[0]}</audio>\n\n{reply}"]
+            else:
+                history['visible'][i] = [history['visible'][i][0], f"{visible_reply.split('</audio>')[0]}</audio>"]
+
+    return history
+
+
+def random_sentence():
+    with open(Path("extensions/coqui_tts/harvard_sentences.txt")) as f:
+        return random.choice(list(f))
+
+
+def voice_preview(string):
+    string = html.unescape(string) or random_sentence()
+
+    output_file = Path('extensions/coqui_tts/outputs/voice_preview.wav')
+    model.tts_to_file(
+        text=string,
+        file_path=output_file,
+        speaker_wav=[f"{this_dir}/voices/{params['voice']}"],
+        language=languages[params["language"]]
+    )
+
+    return f'<audio src="file/{output_file.as_posix()}?{int(time.time())}" controls autoplay></audio>'
+
+
+def history_modifier(history):
+    # Remove autoplay from the last reply
+    if len(history['internal']) > 0:
+        history['visible'][-1] = [
+            history['visible'][-1][0],
+            history['visible'][-1][1].replace('controls autoplay>', 'controls>')
+        ]
+
+    return history
+
+
+def state_modifier(state):
+    if not params['activate']:
+        return state
+
+    state['stream'] = False
+    return state
+
+
+def input_modifier(string, state):
+    if not params['activate']:
+        return string
+
+    shared.processing_message = "*Is recording a voice message...*"
+    return string
+
+
+def output_modifier(string, state):
+    if not params['activate']:
+        return string
+
+    original_string = string
+    string = preprocess(html.unescape(string))
+    if string == '':
+        string = '*Empty reply, try regenerating*'
+    else:
+        output_file = Path(f'extensions/coqui_tts/outputs/{state["character_menu"]}_{int(time.time())}.wav')
+        model.tts_to_file(
+            text=string,
+            file_path=output_file,
+            speaker_wav=[f"{this_dir}/voices/{params['voice']}"],
+            language=languages[params["language"]]
+        )
+
+        autoplay = 'autoplay' if params['autoplay'] else ''
+        string = f'<audio src="file/{output_file.as_posix()}" controls {autoplay}></audio>'
+        if params['show_text']:
+            string += f'\n\n{original_string}'
+
+    shared.processing_message = "*Is typing...*"
+    return string
+
+
+def custom_css():
+    path_to_css = Path(f"{this_dir}/style.css")
+    return open(path_to_css, 'r').read()
+
+
+def setup():
+    global model
+    print("[XTTS] Loading XTTS...")
+    model = load_model()
+    print("[XTTS] Done!")
+    Path(f"{this_dir}/outputs").mkdir(parents=True, exist_ok=True)
+
+
+def ui():
+    with gr.Accordion("Coqui TTS (XTTSv2)"):
+        with gr.Row():
+            activate = gr.Checkbox(value=params['activate'], label='Activate TTS')
+            autoplay = gr.Checkbox(value=params['autoplay'], label='Play TTS automatically')
+
+        with gr.Row():
+            show_text = gr.Checkbox(value=params['show_text'], label='Show message text under audio player')
+            remove_trailing_dots = gr.Checkbox(value=params['remove_trailing_dots'], label='Remove trailing "." from text segments before converting to audio')
+
+        with gr.Row():
+            with gr.Row():
+                voice = gr.Dropdown(get_available_voices(), label="Voice wav", value=params["voice"])
+                create_refresh_button(voice, lambda: None, lambda: {'choices': get_available_voices(), 'value': params["voice"]}, 'refresh-button')
+
+            language = gr.Dropdown(languages.keys(), label="Language", value=params["language"])
+
+        with gr.Row():
+            preview_text = gr.Text(show_label=False, placeholder="Preview text", elem_id="silero_preview_text")
+            preview_play = gr.Button("Preview")
+            preview_audio = gr.HTML(visible=False)
+
+        with gr.Row():
+            convert = gr.Button('Permanently replace audios with the message texts')
+            convert_cancel = gr.Button('Cancel', visible=False)
+            convert_confirm = gr.Button('Confirm (cannot be undone)', variant="stop", visible=False)
+
+    # Convert history with confirmation
+    convert_arr = [convert_confirm, convert, convert_cancel]
+    convert.click(lambda: [gr.update(visible=True), gr.update(visible=False), gr.update(visible=True)], None, convert_arr)
+    convert_confirm.click(
+        lambda: [gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)], None, convert_arr).then(
+        remove_tts_from_history, gradio('history'), gradio('history')).then(
+        chat.save_history, gradio('history', 'unique_id', 'character_menu', 'mode'), None).then(
+        chat.redraw_html, gradio(ui_chat.reload_arr), gradio('display'))
+
+    convert_cancel.click(lambda: [gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)], None, convert_arr)
+
+    # Toggle message text in history
+    show_text.change(
+        lambda x: params.update({"show_text": x}), show_text, None).then(
+        toggle_text_in_history, gradio('history'), gradio('history')).then(
+        chat.save_history, gradio('history', 'unique_id', 'character_menu', 'mode'), None).then(
+        chat.redraw_html, gradio(ui_chat.reload_arr), gradio('display'))
+
+    # Event functions to update the parameters in the backend
+    activate.change(lambda x: params.update({"activate": x}), activate, None)
+    autoplay.change(lambda x: params.update({"autoplay": x}), autoplay, None)
+    remove_trailing_dots.change(lambda x: params.update({"remove_trailing_dots": x}), remove_trailing_dots, None)
+    voice.change(lambda x: params.update({"voice": x}), voice, None)
+    language.change(lambda x: params.update({"language": x}), language, None)
+
+    # Play preview
+    preview_text.submit(voice_preview, preview_text, preview_audio)
+    preview_play.click(voice_preview, preview_text, preview_audio)

extensions/coqui_tts/style.css

@@ -0,0 +1,8 @@
+.SDAP .hires_opts input[type="number"] {
+    width: 6em !important;
+}
+
+/* silero_tts preview */
+.form:has(> #silero_preview_text) {
+    min-width: 75%
+}

extensions/example/script.py

@@ -0,0 +1,139 @@
+"""
+An example of extension. It does nothing, but you can add transformations
+before the return statements to customize the webui behavior.
+
+Starting from history_modifier and ending in output_modifier, the
+functions are declared in the same order that they are called at
+generation time.
+"""
+
+import gradio as gr
+import torch
+from transformers import LogitsProcessor
+
+from modules import chat, shared
+from modules.text_generation import (
+    decode,
+    encode,
+    generate_reply,
+)
+
+params = {
+    "display_name": "Example Extension",
+    "is_tab": False,
+}
+
+class MyLogits(LogitsProcessor):
+    """
+    Manipulates the probabilities for the next token before it gets sampled.
+    Used in the logits_processor_modifier function below.
+    """
+    def __init__(self):
+        pass
+
+    def __call__(self, input_ids, scores):
+        # probs = torch.softmax(scores, dim=-1, dtype=torch.float)
+        # probs[0] /= probs[0].sum()
+        # scores = torch.log(probs / (1 - probs))
+        return scores
+
+def history_modifier(history):
+    """
+    Modifies the chat history.
+    Only used in chat mode.
+    """
+    return history
+
+def state_modifier(state):
+    """
+    Modifies the state variable, which is a dictionary containing the input
+    values in the UI like sliders and checkboxes.
+    """
+    return state
+
+def chat_input_modifier(text, visible_text, state):
+    """
+    Modifies the user input string in chat mode (visible_text).
+    You can also modify the internal representation of the user
+    input (text) to change how it will appear in the prompt.
+    """
+    return text, visible_text
+
+def input_modifier(string, state, is_chat=False):
+    """
+    In default/notebook modes, modifies the whole prompt.
+
+    In chat mode, it is the same as chat_input_modifier but only applied
+    to "text", here called "string", and not to "visible_text".
+    """
+    return string
+
+def bot_prefix_modifier(string, state):
+    """
+    Modifies the prefix for the next bot reply in chat mode.
+    By default, the prefix will be something like "Bot Name:".
+    """
+    return string
+
+def tokenizer_modifier(state, prompt, input_ids, input_embeds):
+    """
+    Modifies the input ids and embeds.
+    Used by the multimodal extension to put image embeddings in the prompt.
+    Only used by loaders that use the transformers library for sampling.
+    """
+    return prompt, input_ids, input_embeds
+
+def logits_processor_modifier(processor_list, input_ids):
+    """
+    Adds logits processors to the list, allowing you to access and modify
+    the next token probabilities.
+    Only used by loaders that use the transformers library for sampling.
+    """
+    processor_list.append(MyLogits())
+    return processor_list
+
+def output_modifier(string, state, is_chat=False):
+    """
+    Modifies the LLM output before it gets presented.
+
+    In chat mode, the modified version goes into history['visible'],
+    and the original version goes into history['internal'].
+    """
+    return string
+
+def custom_generate_chat_prompt(user_input, state, **kwargs):
+    """
+    Replaces the function that generates the prompt from the chat history.
+    Only used in chat mode.
+    """
+    result = chat.generate_chat_prompt(user_input, state, **kwargs)
+    return result
+
+def custom_css():
+    """
+    Returns a CSS string that gets appended to the CSS for the webui.
+    """
+    return ''
+
+def custom_js():
+    """
+    Returns a javascript string that gets appended to the javascript
+    for the webui.
+    """
+    return ''
+
+def setup():
+    """
+    Gets executed only once, when the extension is imported.
+    """
+    pass
+
+def ui():
+    """
+    Gets executed when the UI is drawn. Custom gradio elements and
+    their corresponding event handlers should be defined here.
+
+    To learn about gradio components, check out the docs:
+    https://gradio.app/docs/
+    """
+    pass

extensions/gallery/script.js

@@ -0,0 +1,40 @@
+let gallery_element = document.getElementById('gallery-extension');
+let chat_mode_element = document.getElementById('chat-mode');
+
+let extensions_block = document.getElementById('extensions');
+let extensions_block_size = extensions_block.childNodes.length;
+let gallery_only = (extensions_block_size == 5);
+
+function gotoFirstPage() {
+    const firstPageButton = gallery_element.querySelector('.paginate > button');
+    if (firstPageButton) {
+        firstPageButton.click();
+    }
+}
+
+document.querySelector('.header_bar').addEventListener('click', function(event) {
+    if (event.target.tagName === 'BUTTON') {
+        const buttonText = event.target.textContent.trim();
+
+        let chat_visible = (buttonText == 'Chat');
+        let default_visible = (buttonText == 'Default');
+        let notebook_visible = (buttonText == 'Notebook');
+        let chat_mode_visible = (chat_mode_element.offsetHeight > 0 && chat_mode_element.offsetWidth > 0);
+
+        // Only show this extension in the Chat tab
+        if (chat_visible) {
+            if (chat_mode_visible) {
+                gallery_element.style.display = 'block';
+                extensions_block.style.display = '';
+            } else {
+                gallery_element.style.display = 'none';
+                extensions_block.style.display = 'none';
+            }
+        } else {
+            gallery_element.style.display = 'none';
+            if (gallery_only) {
+                extensions_block.style.display = 'none';
+            }
+        }
+    }
+});

extensions/gallery/script.py

@@ -0,0 +1,129 @@
+from pathlib import Path
+
+import gradio as gr
+
+from modules.html_generator import get_image_cache
+from modules.shared import gradio, settings
+
+
+cards = []
+
+
+def generate_css():
+    css = """
+      .highlighted-border {
+        border-color: rgb(249, 115, 22) !important;
+      }
+
+      .character-gallery > .gallery {
+        margin: 1rem 0;
+        display: grid !important;
+        grid-template-columns: repeat(auto-fit, minmax(150px, 1fr));
+        grid-column-gap: 0.4rem;
+        grid-row-gap: 1.2rem;
+      }
+
+      .character-gallery > .label {
+        display: none !important;
+      }
+
+      .character-gallery button.gallery-item {
+        display: contents;
+      }
+
+      .character-container {
+        cursor: pointer;
+        text-align: center;
+        position: relative;
+        opacity: 0.85;
+      }
+
+      .character-container:hover {
+        opacity: 1;
+      }
+
+      .character-container .placeholder, .character-container img {
+        width: 150px;
+        height: 200px;
+        background-color: gray;
+        object-fit: cover;
+        margin: 0 auto;
+        border-radius: 1rem;
+        border: 3px solid white;
+        box-shadow: 3px 3px 6px 0px rgb(0 0 0 / 50%);
+      }
+
+      .character-name {
+        margin-top: 0.3rem;
+        display: block;
+        font-size: 1.2rem;
+        font-weight: 600;
+        overflow-wrap: anywhere;
+      }
+    """
+    return css
+
+
+def generate_html():
+    global cards
+    cards = []
+    # Iterate through files in image folder
+    for file in sorted(Path("characters").glob("*")):
+        if file.suffix in [".json", ".yml", ".yaml"]:
+            character = file.stem
+            container_html = '<div class="character-container">'
+            image_html = "<div class='placeholder'></div>"
+
+            for path in [Path(f"characters/{character}.{extension}") for extension in ['png', 'jpg', 'jpeg']]:
+                if path.exists():
+                    image_html = f'<img src="file/{get_image_cache(path)}">'
+                    break
+
+            container_html += f'{image_html} <span class="character-name">{character}</span>'
+            container_html += "</div>"
+            cards.append([container_html, character])
+
+    return cards
+
+
+def filter_cards(filter_str=''):
+    if filter_str == '':
+        return cards
+
+    filter_upper = filter_str.upper()
+    return [k for k in cards if filter_upper in k[1].upper()]
+
+
+def select_character(evt: gr.SelectData):
+    return (evt.value[1])
+
+
+def custom_js():
+    path_to_js = Path(__file__).parent.resolve() / 'script.js'
+    return open(path_to_js, 'r').read()
+
+
+def ui():
+    with gr.Accordion("Character gallery", open=settings["gallery-open"], elem_id='gallery-extension'):
+        gr.HTML(value="<style>" + generate_css() + "</style>")
+        with gr.Row():
+            filter_box = gr.Textbox(label='', placeholder='Filter', lines=1, max_lines=1, container=False, elem_id='gallery-filter-box')
+            gr.ClearButton(filter_box, value='Clear', elem_classes='refresh-button')
+            update = gr.Button("Refresh", elem_classes='refresh-button')
+
+        gallery = gr.Dataset(
+            components=[gr.HTML(visible=False)],
+            label="",
+            samples=generate_html(),
+            elem_classes=["character-gallery"],
+            samples_per_page=settings["gallery-items_per_page"]
+        )
+
+    filter_box.change(lambda: None, None, None, _js=f'() => {{{custom_js()}; gotoFirstPage()}}').success(
+        filter_cards, filter_box, gallery).then(
+        lambda x: gr.update(elem_classes='highlighted-border' if x != '' else ''), filter_box, filter_box, show_progress=False)
+
+    update.click(generate_html, [], None).success(
+        filter_cards, filter_box, gallery)
+
+    gallery.select(select_character, None, gradio['character_menu'])

extensions/google_translate/requirements.txt

@@ -0,0 +1 @@
+deep-translator==1.9.2

extensions/google_translate/script.py

@@ -0,0 +1,59 @@
+import html
+
+import gradio as gr
+from deep_translator import GoogleTranslator
+
+params = {
+    "activate": True,
+    "language string": "ja",
+}
+
+language_codes = {'Afrikaans': 'af', 'Albanian': 'sq', 'Amharic': 'am', 'Arabic': 'ar', 'Armenian': 'hy', 'Azerbaijani': 'az', 'Basque': 'eu', 'Belarusian': 'be', 'Bengali': 'bn', 'Bosnian': 'bs', 'Bulgarian': 'bg', 'Catalan': 'ca', 'Cebuano': 'ceb', 'Chinese (Simplified)': 'zh-CN', 'Chinese (Traditional)': 'zh-TW', 'Corsican': 'co', 'Croatian': 'hr', 'Czech': 'cs', 'Danish': 'da', 'Dutch': 'nl', 'English': 'en', 'Esperanto': 'eo', 'Estonian': 'et', 'Finnish': 'fi', 'French': 'fr', 'Frisian': 'fy', 'Galician': 'gl', 'Georgian': 'ka', 'German': 'de', 'Greek': 'el', 'Gujarati': 'gu', 'Haitian Creole': 'ht', 'Hausa': 'ha', 'Hawaiian': 'haw', 'Hebrew': 'iw', 'Hindi': 'hi', 'Hmong': 'hmn', 'Hungarian': 'hu', 'Icelandic': 'is', 'Igbo': 'ig', 'Indonesian': 'id', 'Irish': 'ga', 'Italian': 'it', 'Japanese': 'ja', 'Javanese': 'jw', 'Kannada': 'kn', 'Kazakh': 'kk', 'Khmer': 'km', 'Korean': 'ko', 'Kurdish': 'ku', 'Kyrgyz': 'ky', 'Lao': 'lo', 'Latin': 'la', 'Latvian': 'lv', 'Lithuanian': 'lt', 'Luxembourgish': 'lb', 'Macedonian': 'mk', 'Malagasy': 'mg', 'Malay': 'ms', 'Malayalam': 'ml', 'Maltese': 'mt', 'Maori': 'mi', 'Marathi': 'mr', 'Mongolian': 'mn', 'Myanmar (Burmese)': 'my', 'Nepali': 'ne', 'Norwegian': 'no', 'Nyanja (Chichewa)': 'ny', 'Pashto': 'ps', 'Persian': 'fa', 'Polish': 'pl', 'Portuguese (Portugal, Brazil)': 'pt', 'Punjabi': 'pa', 'Romanian': 'ro', 'Russian': 'ru', 'Samoan': 'sm', 'Scots Gaelic': 'gd', 'Serbian': 'sr', 'Sesotho': 'st', 'Shona': 'sn', 'Sindhi': 'sd', 'Sinhala (Sinhalese)': 'si', 'Slovak': 'sk', 'Slovenian': 'sl', 'Somali': 'so', 'Spanish': 'es', 'Sundanese': 'su', 'Swahili': 'sw', 'Swedish': 'sv', 'Tagalog (Filipino)': 'tl', 'Tajik': 'tg', 'Tamil': 'ta', 'Telugu': 'te', 'Thai': 'th', 'Turkish': 'tr', 'Ukrainian': 'uk', 'Urdu': 'ur', 'Uzbek': 'uz', 'Vietnamese': 'vi', 'Welsh': 'cy', 'Xhosa': 'xh', 'Yiddish': 'yi', 'Yoruba': 'yo', 'Zulu': 'zu'}
+
+
+def input_modifier(string):
+    """
+    This function is applied to your text inputs before
+    they are fed into the model.
+    """
+    if not params['activate']:
+        return string
+
+    return GoogleTranslator(source=params['language string'], target='en').translate(string)
+
+
+def output_modifier(string):
+    """
+    This function is applied to the model outputs.
+    """
+    if not params['activate']:
+        return string
+
+    translated_str = GoogleTranslator(source='en', target=params['language string']).translate(html.unescape(string))
+    return html.escape(translated_str)
+
+
+def bot_prefix_modifier(string):
+    """
+    This function is only applied in chat mode. It modifies
+    the prefix text for the Bot and can be used to bias its
+    behavior.
+    """
+
+    return string
+
+
+def ui():
+    # Finding the language name from the language code to use as the default value
+    language_name = list(language_codes.keys())[list(language_codes.values()).index(params['language string'])]
+
+    # Gradio elements
+    with gr.Row():
+        activate = gr.Checkbox(value=params['activate'], label='Activate translation')
+
+    with gr.Row():
+        language = gr.Dropdown(value=language_name, choices=[k for k in language_codes], label='Language')
+
+    # Event functions to update the parameters in the backend
+    activate.change(lambda x: params.update({"activate": x}), activate, None)
+    language.change(lambda x: params.update({"language string": language_codes[x]}), language, None)

extensions/long_replies/script.py

@@ -0,0 +1,143 @@
+import torch
+from modules import chat, shared
+from modules.text_generation import (
+    decode,
+    encode,
+    generate_reply,
+)
+from transformers import LogitsProcessor
+import gradio as gr
+
+params = {
+    "display_name": "Long replies",
+    "is_tab": False,
+    "min_length": 120,
+}
+
+initial_size = 0
+
+class MyLogits(LogitsProcessor):
+    """
+    Manipulates the probabilities for the next token before it gets sampled.
+    Used in the logits_processor_modifier function below.
+    """
+    def __init__(self):
+        self.newline_id = shared.tokenizer.encode('\n')[-1]
+        pass
+
+    def __call__(self, input_ids, scores):
+        if input_ids.shape[-1] - initial_size < params["min_length"]:
+            scores[...,self.newline_id] = -1000
+            # scores[...,shared.tokenizer.eos_token_id] = -1000
+
+        # probs = torch.softmax(scores, dim=-1, dtype=torch.float)
+        # probs[0] /= probs[0].sum()
+        # scores = torch.log(probs / (1 - probs))
+        return scores
+
+def history_modifier(history):
+    """
+    Modifies the chat history.
+    Only used in chat mode.
+    """
+    return history
+
+def state_modifier(state):
+    """
+    Modifies the state variable, which is a dictionary containing the input
+    values in the UI like sliders and checkboxes.
+    """
+    return state
+
+def chat_input_modifier(text, visible_text, state):
+    """
+    Modifies the user input string in chat mode (visible_text).
+    You can also modify the internal representation of the user
+    input (text) to change how it will appear in the prompt.
+    """
+    return text, visible_text
+
+def input_modifier(string, state):
+    """
+    In default/notebook modes, modifies the whole prompt.
+
+    In chat mode, it is the same as chat_input_modifier but only applied
+    to "text", here called "string", and not to "visible_text".
+    """
+    return string
+
+def bot_prefix_modifier(string, state):
+    """
+    Modifies the prefix for the next bot reply in chat mode.
+    By default, the prefix will be something like "Bot Name:".
+    """
+    return string
+
+def tokenizer_modifier(state, prompt, input_ids, input_embeds):
+    """
+    Modifies the input ids and embeds.
+    Used by the multimodal extension to put image embeddings in the prompt.
+    Only used by loaders that use the transformers library for sampling.
+    """
+
+    global initial_size
+    initial_size = input_ids.shape[-1]
+
+    return prompt, input_ids, input_embeds
+
+def logits_processor_modifier(processor_list, input_ids):
+    """
+    Adds logits processors to the list, allowing you to access and modify
+    the next token probabilities.
+    Only used by loaders that use the transformers library for sampling.
+    """
+    processor_list.append(MyLogits())
+    return processor_list
+
+def output_modifier(string, state):
+    """
+    Modifies the LLM output before it gets presented.
+
+    In chat mode, the modified version goes into history['visible'],
+    and the original version goes into history['internal'].
+    """
+    return string
+
+def custom_generate_chat_prompt(user_input, state, **kwargs):
+    """
+    Replaces the function that generates the prompt from the chat history.
+    Only used in chat mode.
+    """
+    result = chat.generate_chat_prompt(user_input, state, **kwargs)
+    return result
+
+def custom_css():
+    """
+    Returns a CSS string that gets appended to the CSS for the webui.
+    """
+    return ''
+
+def custom_js():
+    """
+    Returns a javascript string that gets appended to the javascript
+    for the webui.
+    """
+    return ''
+
+def setup():
+    """
+    Gets executed only once, when the extension is imported.
+    """
+    pass
+
+def ui():
+    """
+    Gets executed when the UI is drawn. Custom gradio elements and
+    their corresponding event handlers should be defined here.
+
+    To learn about gradio components, check out the docs:
+    https://gradio.app/docs/
+    """
+
+    min_length = gr.Slider(0, 800, step=10, value=params['min_length'], label='Minimum reply length')
+    min_length.change(lambda x: params.update({'min_length': x}), min_length, None)

extensions/multimodal/DOCS.md

@@ -0,0 +1,85 @@
+# Technical description of multimodal extension
+
+## Working principle
+Multimodality extension does most of the stuff which is required for any image input:
+
+- adds the UI
+- saves the images as base64 JPEGs to history
+- provides the hooks to the UI
+- if there are images in the prompt, it:
+    - splits the prompt to text and image parts
+    - adds image start/end markers to text parts, then encodes and embeds the text parts
+    - calls the vision pipeline to embed the images
+    - stitches the embeddings together, and returns them to text generation
+- loads the appropriate vision pipeline, selected either from model name, or by specifying --multimodal-pipeline parameter
+
+Now, for the pipelines, they:
+
+- load the required vision models
+- return some consts, for example the number of tokens taken up by image
+- and most importantly: return the embeddings for LLM, given a list of images
+
+## Prompts/history
+
+To save images in prompt/history, this extension is using a base64 JPEG, wrapped in a HTML tag, like so:
+```
+<img src="data:image/jpeg;base64,{img_str}">
+```
+where `{img_str}` is the actual image data. This format makes displaying them in the UI for free. Do note, that this format is required to be exactly the same, the regex used to find the images is: `<img src="data:image/jpeg;base64,([A-Za-z0-9+/=]+)">`.
+
+## LLM input
+To describe the input, let's see it on an example prompt:
+```
+text1<image1>text2<image2>text3
+```
+where `textN` is N-th text, `<imageN>` is N-th image, in HTML format specified above.
+
+**The first step is to split the prompt into image/text parts**, so we get:
+```
+['text1', '<image1>', 'text2', '<image2>', 'text3']
+```
+this is done in `MultimodalEmbedder._split_prompt(...)` function, which returns a list of `PromptPart`s - dataclasses wrapping the separate parts.
+
+This function also appends the image start/end markers to text, which are provided by `AbstractMultimodalPipeline.image_start()` / `AbstractMultimodalPipeline.image_end()` functions. If image start is `<Img>`, and end is `</Img>`, this function will return:
+```
+['text1<Img>', '<image1>', '</Img>text2<Img>', '<image2>', '</Img>text3']
+```
+
+**The returned prompt parts are then turned into token embeddings.**
+
+First, they are modified to token IDs, for the text it is done using standard `modules.text_generation.encode()` function, and for the images the returned token IDs are changed to placeholders. The placeholder is a list of `N` times `placeholder token id`, where `N` is specified using `AbstractMultimodalPipeline.num_image_embeds()`, and placeholder token IDs using  `AbstractMultimodalPipeline.placeholder_token_id()`.
+
+Now, based on the token IDs, the prompt might get truncated, especially if `max_new_tokens` are unreasonably high. Unfortunately, it can't be done simply, just by trimming the prompt to be short enough. This way will lead to sometimes splitting the prompt in the middle of an image embedding, which usually breaks the generation. Therefore, in this case, the entire image needs to be removed from input. This is done inside `MultimodalEmbedder._encode_text(...)` function.
+
+**After the tokenization, the tokens need to get embedded**, the text and images are once again treated separately.
+
+The text parts are turned to embeddings, using `AbstractMultimodalPipeline.embed_tokens(...)` function. It uses standard embedding function from the model, but to support many LLMs, the actual function is returned by the pipeline (as it might be different for different LLMs), for LLaMA it is `shared.model.model.embed_tokens(...)`.
+
+The image parts are turned to embeddings, using `AbstractMultimodalPipeline.embed_images(...)` function. This function is specific for a given pipeline, it takes the images as input, forwards them through vision model/projector, and returns the embeddings.
+
+**Now, the returned embeddings are stitched together**, using `torch.cat()`, this is creating the final input to the LLM.
+
+## Pipelines
+
+All of the pipelines should subclass `AbstractMultimodalPipeline` class. The idea is to allow for new pipelines to be added in the same way as user extensions - git clone into `extensions/multimodal/pipelines`.
+
+The pipelines are the description of the vision part, containing vision model/multimodal projector. All of the pipelines should have an unique `name()`, which is then selected by user, in `--multimodal-pipeline` CLI argument. For an example, see `pipelines/llava/llava.py`.
+
+## Pipeline modules
+
+Pipelines are organized into "pipeline modules" - subdirectories in `pipelines` directory. The pipeline modules should contain a file called `pipelines.py`, that should contain the following fields:
+- `available_pipelines: List[str]` - list of pipelines provided by this module, shown as the list of available pipelines to the user
+- `def get_pipeline(name: str, params: dict) -> Optional[AbstractMultimodalPipeline]`: - a function to get a concrete pipeline by `name`, if `name` doesn't match any, should return `None`. `params` is the user settings for multimodal extension
+- `def get_pipeline_from_model_name(model_name: str, params: dict) -> Optional[AbstractMultimodalPipeline]`: - a function to get a pipeline from `model_name`, should be eager to return `None`, unless the determination can be done clearly (for example: minigpt-4 bases on vicuna - it should never return the pipeline, but llava can, as it has its own specific LLM finetune)
+
+**NOTE**: A pipeline module should lazy-import the pipelines only when necessary, and it should keep its imports to minimum
+
+## Pipeline params
+
+The pipelines will get the extension `params` in the constructor. They should honor the following fields:
+- `vision_device` - string, specifying `torch.device` to run the vision model (CLIP/ViT) on
+- `vision_bits` - int, number of fp bits to load the vision model(s) in
+- `projector_device` - string, specifying `torch.device` to run the projector models (Linear layers, QFormer, etc.) on
+- `projector_bits` - int, number of fp bits to load the projector models in
+
+As a helper, `AbstractMultimodalPipeline` has `_get_device(self, setting_name: str, params: dict)` and `_get_dtype(self, setting_name: str, params: dict)` helper functions, which parse string/int and return `torch.device` / `torch.dtype`.

extensions/multimodal/README.md

@@ -0,0 +1,139 @@
+# Multimodal
+
+## Description
+
+Adds support for multimodality (text+images) to text-generation-webui.
+
+Note: multimodal currently only works for transformers, AutoGPTQ, and GPTQ-for-LLaMa loaders. ExLlama (v1 and v2) and llama.cpp support are planned.
+
+https://user-images.githubusercontent.com/3718215/233817203-69b57e77-0c55-4fd6-b742-3204bb13b8fc.mp4
+
+## Usage
+
+To run this extension, download a LLM that supports multimodality, and then start server.py with the appropriate `--multimodal-pipeline` argument. Examples:
+
+```
+# LLaVA 1.5 13B has the best performance
+python server.py --model liuhaotian_llava-v1.5-13b --multimodal-pipeline llava-v1.5-13b --load-in-4bit
+# LLaVA 1.5 7B is relatively weaker, but requires less memory
+python server.py --model liuhaotian_llava-v1.5-7b --multimodal-pipeline llava-v1.5-7b --load-in-4bit
+python server.py --model TheBloke_llava-v1.5-13B-GPTQ_gptq-4bit-32g-actorder_True --multimodal-pipeline llava-v1.5-13b --disable_exllama --loader autogptq
+python server.py --model wojtab_llava-7b-v0-4bit-128g --multimodal-pipeline llava-7b
+python server.py --model wojtab_llava-13b-v0-4bit-128g --multimodal-pipeline llava-13b
+python server.py --model anon8231489123_vicuna-13b-GPTQ-4bit-128g --multimodal-pipeline minigpt4-13b
+python server.py --model llama-7b-4bit --multimodal-pipeline minigpt4-7b
+```
+
+There is built-in support for LLaVA-v0-13B, LLaVA-v0-7b, and LLaVA-v1.5-13B. To install `minigpt4`:
+
+- clone https://github.com/Wojtab/minigpt-4-pipeline into `extensions/multimodal/pipelines`
+- install the requirements.txt
+
+The same procedure should be used to install other pipelines, which can then be used with `--multimodal-pipeline [pipeline name]`. For additional multimodal pipelines refer to the compatibility section below.
+
+Do note, that each image takes up a considerable amount of tokens, so adjust `max_new_tokens` to be at most 1700 (recommended value is between 200 to 500), so the images don't get truncated.
+
+To send an image, just upload it to the extension field below chat, and send a prompt as always. The image will be added to the end of your message. If you wish to modify the placement, include a string `<image>` in your prompt.
+
+Additionally, there is *Embed all images, not only the last one* checkbox. It modifies the image embeddings, by default (if it's unchecked), all but the most recent images have their embeddings empty, so they are not fed to the network. It seems as if some multimodal networks consider the features in all images at the same time as if they were a single image. Due to this behavior, by default, the extension skips previous images. However, it can lead to sub-par generation on other pipelines. If you want to include all images, just tick this checkbox.
+
+## Compatibility
+
+As of now, the following multimodal pipelines are supported:
+|Pipeline|`--multimodal-pipeline`|Default LLM|LLM info(for the linked model)|Pipeline repository|
+|-|-|-|-|-|
+|[LLaVA 13B](https://github.com/haotian-liu/LLaVA)|`llava-13b`|[LLaVA 13B](https://huggingface.co/wojtab/llava-13b-v0-4bit-128g)|GPTQ 4-bit quant, old CUDA|built-in|
+|[LLaVA 7B](https://github.com/haotian-liu/LLaVA)|`llava-7b`|[LLaVA 7B](https://huggingface.co/wojtab/llava-7b-v0-4bit-128g)|GPTQ 4-bit quant, old CUDA|built-in|
+|[MiniGPT-4 7B](https://github.com/Vision-CAIR/MiniGPT-4)|`minigpt4-7b`|[Vicuna v0 7B](https://huggingface.co/TheBloke/vicuna-7B-GPTQ-4bit-128g)|GPTQ 4-bit quant, new format|[Wojtab/minigpt-4-pipeline](https://github.com/Wojtab/minigpt-4-pipeline)|
+|[MiniGPT-4 13B](https://github.com/Vision-CAIR/MiniGPT-4)|`minigpt4-13b`|[Vicuna v0 13B](https://huggingface.co/anon8231489123/vicuna-13b-GPTQ-4bit-128g)|GPTQ 4-bit quant, old CUDA|[Wojtab/minigpt-4-pipeline](https://github.com/Wojtab/minigpt-4-pipeline)|
+|[InstructBLIP 7B](https://github.com/salesforce/LAVIS/tree/main/projects/instructblip)|`instructblip-7b`|[Vicuna v1.1 7B](https://huggingface.co/TheBloke/vicuna-7B-1.1-GPTQ-4bit-128g)|GPTQ 4-bit quant|[kjerk/instructblip-pipeline](https://github.com/kjerk/instructblip-pipeline)|
+|[InstructBLIP 13B](https://github.com/salesforce/LAVIS/tree/main/projects/instructblip)|`instructblip-13b`|[Vicuna v1.1 13B](https://huggingface.co/TheBloke/vicuna-13B-1.1-GPTQ-4bit-128g)|GPTQ 4-bit quant|[kjerk/instructblip-pipeline](https://github.com/kjerk/instructblip-pipeline)|
+
+Some pipelines could support different LLMs but do note that while it might work, it isn't a supported configuration.
+
+DO NOT report bugs if you are using a different LLM.
+
+DO NOT report bugs with pipelines in this repository (unless they are built-in)
+
+## Extension config
+This extension uses the following parameters (from `settings.json`):
+|Parameter|Description|
+|---------|-----------|
+|`multimodal-vision_bits`|Number of bits to load vision models (CLIP/ViT) feature extractor in (most pipelines should support either 32 or 16, default=32)|
+|`multimodal-vision_device`|Torch device to run the feature extractor on, for example, `cpu` or `cuda:0`, by default `cuda:0` if available|
+|`multimodal-projector_bits`|Number of bits to load feature projector model(s) in (most pipelines should support either 32 or 16, default=32)|
+|`multimodal-projector_device`|Torch device to run the feature projector model(s) on, for example `cpu` or `cuda:0`, by default `cuda:0` if available|
+|`multimodal-add_all_images_to_prompt`|Default value of "Embed all images, not only the last one" checkbox|
+
+## Usage through API
+
+### Chat completions endpoint
+
+#### With an image URL
+
+```shell
+curl http://127.0.0.1:5000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      {
+        "role": "user",
+        "image_url": "https://avatars.githubusercontent.com/u/112222186?v=4"
+      },
+      {
+        "role": "user",
+        "content": "What is unusual about this image?"
+      }
+    ]
+  }'
+```
+
+#### With a Base64 image
+
+```python
+import base64
+import json
+import requests
+
+img = open('image.jpg', 'rb')
+img_bytes = img.read()
+img_base64 = base64.b64encode(img_bytes).decode('utf-8')
+data = { "messages": [
+        {
+            "role": "user",
+            "image_url": f"data:image/jpeg;base64,{img_base64}"
+        },
+        {
+            "role": "user",
+            "content": "what is unusual about this image?"
+        }
+    ]
+}
+response = requests.post('http://127.0.0.1:5000/v1/chat/completions', json=data)
+print(response.text)
+```
+
+You can run the multimodal inference through API, by inputting the images to prompt. Images are embedded like so: `f'<img src="data:image/jpeg;base64,{img_str}">'`, where `img_str` is base-64 jpeg data. Note that you will need to launch `server.py` with the arguments `--api --extensions multimodal`. 
+
+### Completions endpoint
+
+Python example:
+
+```Python
+import base64
+import requests
+
+CONTEXT = "You are LLaVA, a large language and vision assistant trained by UW Madison WAIV Lab. You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language. Follow the instructions carefully and explain your answers in detail.### Human: Hi!### Assistant: Hi there! How can I help you today?\n"
+
+with open('extreme_ironing.jpg', 'rb') as f:
+    img_str = base64.b64encode(f.read()).decode('utf-8')
+    prompt = CONTEXT + f'### Human: What is unusual about this image: \n<img src="data:image/jpeg;base64,{img_str}">### Assistant: '
+    print(requests.post('http://127.0.0.1:5000/v1/completions', json={'prompt': prompt, 'max_tokens': 200, 'stop': ['\n###']}).json())
+```
+script output:
+```Python
+{'results': [{'text': "The unusual aspect of this image is that a man is standing on top of a yellow minivan while doing his laundry. He has set up a makeshift clothes line using the car's rooftop as an outdoor drying area. This scene is uncommon because people typically do their laundry indoors, in a dedicated space like a laundromat or a room in their home, rather than on top of a moving vehicle. Additionally, hanging clothes on the car could be potentially hazardous or illegal in some jurisdictions due to the risk of damaging the vehicle or causing accidents on the road.\n##"}]}
+```
+
+## For pipeline developers/technical description
+see [DOCS.md](https://github.com/oobabooga/text-generation-webui/blob/main/extensions/multimodal/DOCS.md)

extensions/multimodal/abstract_pipeline.py

@@ -0,0 +1,63 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+import torch
+from PIL import Image
+from transformers import is_torch_xpu_available
+
+
+class AbstractMultimodalPipeline(ABC):
+    @staticmethod
+    @abstractmethod
+    def name() -> str:
+        'name of the pipeline, should be same as in --multimodal-pipeline'
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def image_start() -> Optional[str]:
+        'return image start string, string representation of image start token, or None if not applicable'
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def image_end() -> Optional[str]:
+        'return image end string, string representation of image end token, or None if not applicable'
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def placeholder_token_id() -> int:
+        'return placeholder token id'
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def num_image_embeds() -> int:
+        'return the number of embeds used by a single image (for example: 256 for LLaVA)'
+        pass
+
+    @abstractmethod
+    def embed_images(self, images: List[Image.Image]) -> torch.Tensor:
+        'forward the images through vision pipeline, and return their embeddings'
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def embed_tokens(input_ids: torch.Tensor) -> torch.Tensor:
+        'embed tokens, the exact function varies by LLM, for LLaMA it is `shared.model.model.embed_tokens`'
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def placeholder_embeddings() -> torch.Tensor:
+        'get placeholder embeddings if there are multiple images, and `add_all_images_to_prompt` is False'
+        pass
+
+    def _get_device(self, setting_name: str, params: dict):
+        if params[setting_name] is None:
+            return torch.device("cuda:0" if torch.cuda.is_available() else "xpu:0" if is_torch_xpu_available() else "cpu")
+        return torch.device(params[setting_name])
+
+    def _get_dtype(self, setting_name: str, params: dict):
+        return torch.float32 if int(params[setting_name]) == 32 else torch.float16

extensions/multimodal/multimodal_embedder.py

@@ -0,0 +1,178 @@
+import base64
+import re
+from dataclasses import dataclass
+from io import BytesIO
+from typing import Any, List, Optional
+
+import torch
+from PIL import Image
+
+from extensions.multimodal.pipeline_loader import load_pipeline
+from modules import shared
+from modules.logging_colors import logger
+from modules.text_generation import encode, get_max_prompt_length
+
+
+@dataclass
+class PromptPart:
+    text: str
+    image: Optional[Image.Image] = None
+    is_image: bool = False
+    input_ids: Optional[torch.Tensor] = None
+    embedding: Optional[torch.Tensor] = None
+
+
+class MultimodalEmbedder:
+    def __init__(self, params: dict):
+        pipeline, source = load_pipeline(params)
+        self.pipeline = pipeline
+        logger.info(f'Multimodal: loaded pipeline {self.pipeline.name()} from pipelines/{source} ({self.pipeline.__class__.__name__})')
+
+    def _split_prompt(self, prompt: str, load_images: bool = False) -> List[PromptPart]:
+        """Splits a prompt into a list of `PromptParts` to separate image data from text.
+        It will also append `image_start` and `image_end` before and after the image, and optionally parse and load the images,
+        if `load_images` is `True`.
+        """
+        parts: List[PromptPart] = []
+        curr = 0
+        while True:
+            match = re.search(r'<img src="data:image/jpeg;base64,([A-Za-z0-9+/=]+)">', prompt[curr:])
+            if match is None:
+                # no more image tokens, append the rest of the prompt
+                if curr > 0:
+                    # add image end token after last image
+                    parts.append(PromptPart(text=self.pipeline.image_end() + prompt[curr:]))
+                else:
+                    parts.append(PromptPart(text=prompt))
+                break
+            # found an image, append image start token to the text
+            if match.start() > 0:
+                parts.append(PromptPart(text=prompt[curr:curr + match.start()] + self.pipeline.image_start()))
+            else:
+                parts.append(PromptPart(text=self.pipeline.image_start()))
+            # append the image
+            parts.append(PromptPart(
+                text=match.group(0),
+                image=Image.open(BytesIO(base64.b64decode(match.group(1)))) if load_images else None,
+                is_image=True
+            ))
+            curr += match.end()
+        return parts
+
+    def _len_in_tokens_prompt_parts(self, parts: List[PromptPart]) -> int:
+        """Total length in tokens of all `parts`"""
+        tokens = 0
+        for part in parts:
+            if part.is_image:
+                tokens += self.pipeline.num_image_embeds()
+            elif part.input_ids is not None:
+                tokens += len(part.input_ids)
+            else:
+                tokens += len(encode(part.text)[0])
+        return tokens
+
+    def len_in_tokens(self, prompt: str) -> int:
+        """Total length in tokens for a given text `prompt`"""
+        parts = self._split_prompt(prompt, False)
+        return self._len_in_tokens_prompt_parts(parts)
+
+    def _encode_single_text(self, part: PromptPart, add_bos_token: bool) -> PromptPart:
+        """Encode a single prompt `part` to `input_ids`. Returns a `PromptPart`"""
+        if part.is_image:
+            placeholders = torch.ones((self.pipeline.num_image_embeds())) * self.pipeline.placeholder_token_id()
+            part.input_ids = placeholders.to(shared.model.device, dtype=torch.int64)
+        else:
+            part.input_ids = encode(part.text, add_bos_token=add_bos_token)[0].to(shared.model.device, dtype=torch.int64)
+        return part
+
+    @staticmethod
+    def _num_images(parts: List[PromptPart]) -> int:
+        count = 0
+        for part in parts:
+            if part.is_image:
+                count += 1
+        return count
+
+    def _encode_text(self, state, parts: List[PromptPart]) -> List[PromptPart]:
+        """Encode text to token_ids, also truncate the prompt, if necessary.
+
+        The chat/instruct mode should make prompts that fit in get_max_prompt_length, but if max_new_tokens are set
+        such that the context + min_rows don't fit, we can get a prompt which is too long.
+        We can't truncate image embeddings, as it leads to broken generation, so remove the images instead and warn the user
+        """
+        encoded: List[PromptPart] = []
+        for i, part in enumerate(parts):
+            encoded.append(self._encode_single_text(part, i == 0 and state['add_bos_token']))
+
+        # truncation:
+        max_len = get_max_prompt_length(state)
+        removed_images = 0
+
+        # 1. remove entire text/image blocks
+        while self._len_in_tokens_prompt_parts(encoded[1:]) > max_len:
+            if encoded[0].is_image:
+                removed_images += 1
+            encoded = encoded[1:]
+
+        # 2. check if the last prompt part doesn't need to get truncated
+        if self._len_in_tokens_prompt_parts(encoded) > max_len:
+            if encoded[0].is_image:
+                # don't truncate image embeddings, just remove the image, otherwise generation will be broken
+                removed_images += 1
+                encoded = encoded[1:]
+            elif len(encoded) > 1 and encoded[0].text.endswith(self.pipeline.image_start()):
+                # see if we can keep image_start token
+                len_image_start = len(encode(self.pipeline.image_start(), add_bos_token=state['add_bos_token'])[0])
+                if self._len_in_tokens_prompt_parts(encoded[1:]) + len_image_start > max_len:
+                    # we can't -> remove this text, and the image
+                    encoded = encoded[2:]
+                    removed_images += 1
+                else:
+                    # we can -> just truncate the text
+                    trunc_len = self._len_in_tokens_prompt_parts(encoded) - max_len
+                    encoded[0].input_ids = encoded[0].input_ids[trunc_len:]
+            elif len(encoded) > 0:
+                # only one text left, truncate it normally
+                trunc_len = self._len_in_tokens_prompt_parts(encoded) - max_len
+                encoded[0].input_ids = encoded[0].input_ids[trunc_len:]
+
+        # notify user if we truncated an image
+        if removed_images > 0:
+            logger.warning(f"Multimodal: removed {removed_images} image(s) from prompt. Try decreasing max_new_tokens if generation is broken")
+
+        return encoded
+
+    def _embed(self, parts: List[PromptPart]) -> List[PromptPart]:
+        # batch images
+        image_indicies = [i for i, part in enumerate(parts) if part.is_image]
+        embedded = self.pipeline.embed_images([parts[i].image for i in image_indicies])
+        for i, embeds in zip(image_indicies, embedded):
+            parts[i].embedding = embeds
+        # embed text
+        for (i, part) in enumerate(parts):
+            if not part.is_image:
+                parts[i].embedding = self.pipeline.embed_tokens(part.input_ids)
+        return parts
+
+    def _remove_old_images(self, parts: List[PromptPart], params: dict) -> List[PromptPart]:
+        if params['add_all_images_to_prompt']:
+            return parts
+        already_added = False
+        for i, part in reversed(list(enumerate(parts))):
+            if part.is_image:
+                if already_added:
+                    parts[i].embedding = self.pipeline.placeholder_embeddings()
+                else:
+                    already_added = True
+        return parts
+
+    def forward(self, prompt: str, state: Any, params: dict):
+        prompt_parts = self._split_prompt(prompt, True)
+        prompt_parts = self._encode_text(state, prompt_parts)
+        prompt_parts = self._embed(prompt_parts)
+        prompt_parts = self._remove_old_images(prompt_parts, params)
+        embeds = tuple(part.embedding for part in prompt_parts)
+        ids = tuple(part.input_ids for part in prompt_parts)
+        input_embeds = torch.cat(embeds, dim=0)
+        input_ids = torch.cat(ids, dim=0)
+        return prompt, input_ids, input_embeds, self._num_images(prompt_parts)

extensions/multimodal/pipeline_loader.py

@@ -0,0 +1,52 @@
+import traceback
+from importlib import import_module
+from pathlib import Path
+from typing import Tuple
+
+from extensions.multimodal.abstract_pipeline import AbstractMultimodalPipeline
+from modules import shared
+from modules.logging_colors import logger
+
+
+def _get_available_pipeline_modules():
+    pipeline_path = Path(__file__).parent / 'pipelines'
+    modules = [p for p in pipeline_path.iterdir() if p.is_dir()]
+    return [m.name for m in modules if (m / 'pipelines.py').exists()]
+
+
+def load_pipeline(params: dict) -> Tuple[AbstractMultimodalPipeline, str]:
+    pipeline_modules = {}
+    available_pipeline_modules = _get_available_pipeline_modules()
+    for name in available_pipeline_modules:
+        try:
+            pipeline_modules[name] = import_module(f'extensions.multimodal.pipelines.{name}.pipelines')
+        except:
+            logger.warning(f'Failed to get multimodal pipelines from {name}')
+            logger.warning(traceback.format_exc())
+
+    if shared.args.multimodal_pipeline is not None:
+        for k in pipeline_modules:
+            if hasattr(pipeline_modules[k], 'get_pipeline'):
+                pipeline = getattr(pipeline_modules[k], 'get_pipeline')(shared.args.multimodal_pipeline, params)
+                if pipeline is not None:
+                    return (pipeline, k)
+    else:
+        model_name = shared.args.model.lower()
+        for k in pipeline_modules:
+            if hasattr(pipeline_modules[k], 'get_pipeline_from_model_name'):
+                pipeline = getattr(pipeline_modules[k], 'get_pipeline_from_model_name')(model_name, params)
+                if pipeline is not None:
+                    return (pipeline, k)
+
+    available = []
+    for k in pipeline_modules:
+        if hasattr(pipeline_modules[k], 'available_pipelines'):
+            pipelines = getattr(pipeline_modules[k], 'available_pipelines')
+            available += pipelines
+
+    if shared.args.multimodal_pipeline is not None:
+        log = f'Multimodal - ERROR: Failed to load multimodal pipeline "{shared.args.multimodal_pipeline}", available pipelines are: {available}.'
+    else:
+        log = f'Multimodal - ERROR: Failed to determine multimodal pipeline for model {shared.args.model}, please select one manually using --multimodal-pipeline [PIPELINE]. Available pipelines are: {available}.'
+    logger.critical(f'{log} Please specify a correct pipeline, or disable the extension')
+    raise RuntimeError(f'{log} Please specify a correct pipeline, or disable the extension')

extensions/multimodal/pipelines/llava/README.md

@@ -0,0 +1,9 @@
+## LLaVA pipeline
+
+This module provides 2 pipelines:
+- `llava-7b` - for use with LLaVA v0 7B model (finetuned LLaMa 7B)
+- `llava-13b` - for use with LLaVA v0 13B model (finetuned LLaMa 13B)
+
+[LLaVA](https://github.com/haotian-liu/LLaVA) uses CLIP `openai/clip-vit-large-patch14` as the vision model, and then a single linear layer. For 13B the projector weights are in `liuhaotian/LLaVA-13b-delta-v0`, and for 7B they are in `liuhaotian/LLaVA-7b-delta-v0`.
+
+The supported parameter combinations for both the vision model, and the projector are: CUDA/32bit, CUDA/16bit, CPU/32bit

extensions/multimodal/pipelines/llava/llava.py

@@ -0,0 +1,262 @@
+import time
+from abc import abstractmethod
+from typing import List, Tuple
+
+import torch
+from huggingface_hub import hf_hub_download
+from PIL import Image
+from transformers import CLIPImageProcessor, CLIPVisionModel
+
+from extensions.multimodal.abstract_pipeline import AbstractMultimodalPipeline
+from modules import shared
+from modules.logging_colors import logger
+from modules.text_generation import encode
+
+
+def expand2square(pil_img: Image.Image, background_color: Tuple[int]) -> Image.Image:
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+
+class LLaVA_v0_Pipeline(AbstractMultimodalPipeline):
+    CLIP_REPO = "openai/clip-vit-large-patch14"
+
+    def __init__(self, params: dict) -> None:
+        super().__init__()
+        self.clip_device = self._get_device("vision_device", params)
+        self.clip_dtype = self._get_dtype("vision_bits", params)
+        self.projector_device = self._get_device("projector_device", params)
+        self.projector_dtype = self._get_dtype("projector_bits", params)
+        self.image_processor, self.vision_tower, self.mm_projector = self._load_models()
+
+    def _load_models(self):
+        start_ts = time.time()
+
+        logger.info(f"LLaVA - Loading CLIP from {self.CLIP_REPO} as {self.clip_dtype} on {self.clip_device}...")
+        image_processor = CLIPImageProcessor.from_pretrained(self.CLIP_REPO, torch_dtype=self.clip_dtype)
+        vision_tower = CLIPVisionModel.from_pretrained(self.CLIP_REPO, torch_dtype=self.clip_dtype).to(self.clip_device)
+
+        logger.info(f"LLaVA - Loading projector from {self.llava_projector_repo()} as {self.projector_dtype} on {self.projector_device}...")
+        projector_path = hf_hub_download(self.llava_projector_repo(), self.llava_projector_filename())
+        mm_projector = self.build_mm_projector()
+        projector_data = torch.load(projector_path)
+        projector_data = {k[19:]: v for k, v in projector_data.items() if k.startswith('model.mm_projector.')}
+        mm_projector.load_state_dict(projector_data)
+        mm_projector = mm_projector.to(self.projector_device)
+
+        logger.info(f"LLaVA supporting models loaded, took {time.time() - start_ts:.2f} seconds")
+        return image_processor, vision_tower, mm_projector
+
+    def build_mm_projector(self) -> torch.nn.Module:
+        projector_shape = self.llava_projector_shape()
+        if len(projector_shape) == 2:
+            return torch.nn.Linear(*projector_shape)
+        else:
+            modules = []
+            modules.append(torch.nn.Linear(projector_shape[0], projector_shape[1]))
+            for i in range(2, len(projector_shape)):
+                modules.append(torch.nn.GELU())
+                modules.append(torch.nn.Linear(projector_shape[i-1], projector_shape[i]))
+            return torch.nn.Sequential(*modules)
+
+    @staticmethod
+    def image_start() -> str:
+        return "<im_start>"
+
+    @staticmethod
+    def image_end() -> str:
+        return "<im_end>"
+
+    @staticmethod
+    def num_image_embeds() -> int:
+        return 256
+
+    @staticmethod
+    def embed_tokens(input_ids: torch.Tensor) -> torch.Tensor:
+        for attr in ['', 'model', 'model.model', 'model.model.model']:
+            tmp = getattr(shared.model, attr, None) if attr != '' else shared.model
+            if tmp is not None and hasattr(tmp, 'embed_tokens'):
+                func = tmp.embed_tokens
+                break
+        else:
+            raise ValueError('The embed_tokens method has not been found for this loader.')
+
+        return func(input_ids).to(shared.model.device, dtype=shared.model.dtype)
+
+    @staticmethod
+    def placeholder_embeddings() -> torch.Tensor:
+        return LLaVA_v0_Pipeline.embed_tokens(encode("<im_patch>"*256, add_bos_token=False)[0])
+
+    def embed_images(self, images: List[Image.Image]) -> torch.Tensor:
+        images = self.image_processor(images, return_tensors='pt')['pixel_values']
+        images = images.to(self.clip_device, dtype=self.clip_dtype)
+
+        with torch.no_grad():
+            image_forward_outs = self.vision_tower(images, output_hidden_states=True)
+            select_hidden_state_layer = -2
+            select_hidden_state = image_forward_outs.hidden_states[select_hidden_state_layer]
+            image_features = select_hidden_state[:, 1:].to(self.projector_device, dtype=self.projector_dtype)
+            image_features = self.mm_projector(image_features)
+        return image_features.to(shared.model.device, dtype=shared.model.dtype)
+
+    @staticmethod
+    @abstractmethod
+    def llava_projector_repo() -> str:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def llava_projector_filename() -> str:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def llava_projector_shape() -> Tuple[int, int]:
+        pass
+
+
+class LLaVA_v0_13B_Pipeline(LLaVA_v0_Pipeline):
+    def __init__(self, params: dict) -> None:
+        super().__init__(params)
+
+    @staticmethod
+    def name() -> str:
+        return "llava-13b"
+
+    @staticmethod
+    def placeholder_token_id() -> int:
+        return 32000
+
+    @staticmethod
+    def llava_projector_shape() -> Tuple[int, int]:
+        return (1024, 5120)
+
+    @staticmethod
+    def llava_projector_filename() -> str:
+        return "mm_projector.bin"
+
+    @staticmethod
+    def llava_projector_repo() -> str:
+        return "liuhaotian/LLaVA-13b-delta-v0"
+
+
+class LLaVA_v0_7B_Pipeline(LLaVA_v0_Pipeline):
+    def __init__(self, params: dict) -> None:
+        super().__init__(params)
+
+    @staticmethod
+    def name() -> str:
+        return "llava-7b"
+
+    @staticmethod
+    def placeholder_token_id() -> int:
+        return 32001
+
+    @staticmethod
+    def llava_projector_shape() -> Tuple[int, int]:
+        return (1024, 4096)
+
+    @staticmethod
+    def llava_projector_filename() -> str:
+        return "mm_projector.bin"
+
+    @staticmethod
+    def llava_projector_repo() -> str:
+        return "liuhaotian/LLaVA-7b-delta-v0"
+
+
+class LLaVA_LLaMA_2_13B_Pipeline(LLaVA_v0_13B_Pipeline):
+    def __init__(self, params: dict) -> None:
+        super().__init__(params)
+
+    @staticmethod
+    def name() -> str:
+        return "llava-llama-2-13b"
+
+    @staticmethod
+    def placeholder_token_id() -> int:
+        return 0
+
+    @staticmethod
+    def llava_projector_repo() -> str:
+        return "liuhaotian/llava-llama-2-13b-chat-lightning-preview"
+
+    @staticmethod
+    def image_start() -> str:
+        return ""
+
+    @staticmethod
+    def image_end() -> str:
+        return ""
+
+    @staticmethod
+    def placeholder_embeddings() -> torch.Tensor:
+        return LLaVA_v0_Pipeline.embed_tokens(encode("<unk>"*256, add_bos_token=False)[0])
+
+
+class LLaVA_v1_5_13B_Pipeline(LLaVA_v0_13B_Pipeline):
+    CLIP_REPO = "openai/clip-vit-large-patch14-336"
+
+    def __init__(self, params: dict) -> None:
+        super().__init__(params)
+
+    @staticmethod
+    def name() -> str:
+        return "llava-v1.5-13b"
+
+    @staticmethod
+    def llava_projector_shape() -> Tuple[int, int]:
+        return (1024, 5120, 5120)
+
+    @staticmethod
+    def placeholder_token_id() -> int:
+        return 0
+
+    @staticmethod
+    def llava_projector_repo() -> str:
+        return "liuhaotian/llava-v1.5-13b"
+
+    @staticmethod
+    def image_start() -> str:
+        return ""
+
+    @staticmethod
+    def image_end() -> str:
+        return ""
+
+    @staticmethod
+    def num_image_embeds() -> int:
+        return 576
+
+    def embed_images(self, images: List[Image.Image]) -> torch.Tensor:
+        # pad it to square first
+        images = [
+            expand2square(image, tuple(int(x*255) for x in self.image_processor.image_mean))
+            for image in images
+        ]
+        return super().embed_images(images)
+
+    @staticmethod
+    def placeholder_embeddings() -> torch.Tensor:
+        return LLaVA_v0_Pipeline.embed_tokens(encode("<unk>"*576, add_bos_token=False)[0])
+
+class LLaVA_v1_5_7B_Pipeline(LLaVA_v1_5_13B_Pipeline):
+    @staticmethod
+    def name() -> str:
+        return "llava-v1.5-7b"
+
+    @staticmethod
+    def llava_projector_shape() -> Tuple[int, int]:
+        return (1024, 4096, 4096)
+    @staticmethod
+    def llava_projector_repo() -> str:
+        return "liuhaotian/llava-v1.5-7b"

extensions/multimodal/pipelines/llava/pipelines.py

@@ -0,0 +1,48 @@
+from typing import Optional
+
+from extensions.multimodal.abstract_pipeline import AbstractMultimodalPipeline
+
+available_pipelines = ['llava-7b', 'llava-13b', 'llava-llama-2-13b', 'llava-v1.5-13b', 'llava-v1.5-7b']
+
+
+def get_pipeline(name: str, params: dict) -> Optional[AbstractMultimodalPipeline]:
+    if name == 'llava-7b':
+        from .llava import LLaVA_v0_7B_Pipeline
+        return LLaVA_v0_7B_Pipeline(params)
+    if name == 'llava-13b':
+        from .llava import LLaVA_v0_13B_Pipeline
+        return LLaVA_v0_13B_Pipeline(params)
+    if name == 'llava-llama-2-13b':
+        from .llava import LLaVA_LLaMA_2_13B_Pipeline
+        return LLaVA_LLaMA_2_13B_Pipeline(params)
+    if name == 'llava-v1.5-7b':
+        from .llava import LLaVA_v1_5_7B_Pipeline
+        return LLaVA_v1_5_7B_Pipeline(params)
+    if name == 'llava-v1.5-13b':
+        from .llava import LLaVA_v1_5_13B_Pipeline
+        return LLaVA_v1_5_13B_Pipeline(params)
+    return None
+
+
+def get_pipeline_from_model_name(model_name: str, params: dict) -> Optional[AbstractMultimodalPipeline]:
+    if 'llava' not in model_name.lower():
+        return None
+    if 'llama-2' in model_name.lower():
+        if '13b' in model_name.lower():
+            from .llava import LLaVA_LLaMA_2_13B_Pipeline
+            return LLaVA_LLaMA_2_13B_Pipeline(params)
+    elif 'llava-v1.5' in model_name.lower():
+        if '13b' in model_name.lower():
+            from .llava import LLaVA_v1_5_13B_Pipeline
+            return LLaVA_v1_5_13B_Pipeline(params)
+        if '7b' in model_name.lower():
+            from .llava import LLaVA_v1_5_7B_Pipeline
+            return LLaVA_v1_5_7B_Pipeline(params)
+    else:
+        if '7b' in model_name.lower():
+            from .llava import LLaVA_v0_7B_Pipeline
+            return LLaVA_v0_7B_Pipeline(params)
+        if '13b' in model_name.lower():
+            from .llava import LLaVA_v0_13B_Pipeline
+            return LLaVA_v0_13B_Pipeline(params)
+    return None

extensions/multimodal/script.py

@@ -0,0 +1,113 @@
+import base64
+import re
+import time
+from functools import partial
+from io import BytesIO
+
+import gradio as gr
+import torch
+
+from extensions.multimodal.multimodal_embedder import MultimodalEmbedder
+from modules import shared
+from modules.logging_colors import logger
+
+params = {
+    "add_all_images_to_prompt": False,
+    # device to run vision encoder on
+    "vision_device": None,
+    # bits to load vision encoder in, either 16 or 32
+    "vision_bits": 32,
+    # device to run multimodal projector on
+    "projector_device": None,
+    # multimodal projector bits, either 32 or 16
+    "projector_bits": 32
+}
+
+
+# If 'state' is True, will hijack the next chat generation
+input_hijack = {
+    'state': False,
+    'value': ["", ""]
+}
+
+
+# initialized in ui, so that params are loaded from settings
+multimodal_embedder: MultimodalEmbedder = None
+
+
+def chat_input_modifier(text, visible_text, state):
+    global input_hijack
+    if input_hijack['state']:
+        input_hijack['state'] = False
+        return input_hijack['value'](text, visible_text)
+    else:
+        return text, visible_text
+
+
+def add_chat_picture(picture, text, visible_text):
+    # resize the image, so that shortest edge is at least 224 (size for CLIP), and at most 300 (to keep history manageable)
+    # Adjusted to 336 for the values here, due to the increased resolution in llava-v1.5
+    max_hw, min_hw = max(picture.size), min(picture.size)
+    aspect_ratio = max_hw / min_hw
+    shortest_edge = int(max(336 / aspect_ratio, 336))
+    longest_edge = int(shortest_edge * aspect_ratio)
+    w = shortest_edge if picture.width < picture.height else longest_edge
+    h = shortest_edge if picture.width >= picture.height else longest_edge
+    picture = picture.resize((w, h))
+
+    buffer = BytesIO()
+    picture.save(buffer, format="PNG")
+    img_str = base64.b64encode(buffer.getvalue()).decode('utf-8')
+    image = f'<img src="data:image/jpeg;base64,{img_str}">'
+
+    if '<image>' in text:
+        text = text.replace('<image>', image)
+    else:
+        text = image + '\n' + text
+
+    if visible_text == '' or visible_text is None:
+        visible_text = text
+    elif '<image>' in visible_text:
+        visible_text = visible_text.replace('<image>', image)
+    else:
+        visible_text = visible_text + '\n' + image
+
+    return text, visible_text
+
+
+def custom_tokenized_length(prompt):
+    return multimodal_embedder.len_in_tokens(prompt)
+
+
+def tokenizer_modifier(state, prompt, input_ids, input_embeds):
+    global params
+    start_ts = time.time()
+    image_match = re.search(r'<img src="data:image/jpeg;base64,[A-Za-z0-9+/=]+">', prompt)
+
+    if image_match is None:
+        return prompt, input_ids, input_embeds
+
+    prompt, input_ids, input_embeds, total_embedded = multimodal_embedder.forward(prompt, state, params)
+    logger.info(f'Embedded {total_embedded} image(s) in {time.time()-start_ts:.2f}s')
+    return (prompt,
+            input_ids.unsqueeze(0).to(shared.model.device, dtype=torch.int64),
+            input_embeds.unsqueeze(0).to(shared.model.device, dtype=shared.model.dtype))
+
+
+def ui():
+    global multimodal_embedder
+    multimodal_embedder = MultimodalEmbedder(params)
+    with gr.Column():
+        picture_select = gr.Image(label='Send a picture', type='pil')
+        # The models don't seem to deal well with multiple images
+        single_image_checkbox = gr.Checkbox(False, label='Embed all images, not only the last one')
+    # Prepare the input hijack
+    picture_select.upload(
+        lambda picture: input_hijack.update({"state": True, "value": partial(add_chat_picture, picture)}),
+        [picture_select],
+        None
+    )
+    picture_select.clear(lambda: input_hijack.update({"state": False, "value": ["", ""]}), None, None)
+    single_image_checkbox.change(lambda x: params.update({"add_all_images_to_prompt": x}), single_image_checkbox, None)
+    shared.gradio['Generate'].click(lambda: None, None, picture_select)
+    shared.gradio['textbox'].submit(lambda: None, None, picture_select)

extensions/ngrok/README.md

@@ -0,0 +1,69 @@
+# Adding an ingress URL through the ngrok Agent SDK for Python
+
+[ngrok](https://ngrok.com) is a globally distributed reverse proxy commonly used for quickly getting a public URL to a
+service running inside a private network, such as on your local laptop. The ngrok agent is usually
+deployed inside a private network and is used to communicate with the ngrok cloud service.
+
+By default the authtoken in the NGROK_AUTHTOKEN environment variable will be used. Alternatively one may be specified in
+the `settings.json` file, see the Examples below. Retrieve your authtoken on the [Auth Token page of your ngrok dashboard](https://dashboard.ngrok.com/get-started/your-authtoken), signing up is free.
+
+# Documentation
+
+For a list of all available options, see [the configuration documentation](https://ngrok.com/docs/ngrok-agent/config/) or [the connect example](https://github.com/ngrok/ngrok-py/blob/main/examples/ngrok-connect-full.py).
+
+The ngrok Python SDK is [on github here](https://github.com/ngrok/ngrok-py). A quickstart guide and a full API reference are included in the [ngrok-py Python API documentation](https://ngrok.github.io/ngrok-py/).
+
+# Running
+
+To enable ngrok install the requirements and then add `--extension ngrok` to the command line options, for instance:
+
+```bash
+pip install -r extensions/ngrok/requirements.txt
+python server.py --extension ngrok
+```
+
+In the output you should then see something like this:
+
+```bash
+INFO:Loading the extension "ngrok"...
+INFO:Session created
+INFO:Created tunnel "9d9d0944dc75ff9d3aae653e5eb29fe9" with url "https://d83706cf7be7.ngrok.app"
+INFO:Tunnel "9d9d0944dc75ff9d3aae653e5eb29fe9" TCP forwarding to "localhost:7860"
+INFO:Ingress established at https://d83706cf7be7.ngrok.app
+```
+
+You can now access the webui via the url shown, in this case `https://d83706cf7be7.ngrok.app`. It is recommended to add some authentication to the ingress, see below.
+
+# Example Settings
+
+In `settings.json` add a `ngrok` key with a dictionary of options, for instance:
+
+To enable basic authentication:
+```json
+{
+    "ngrok": {
+        "basic_auth": "user:password"
+    }
+}
+```
+
+To enable OAUTH authentication:
+```json
+{
+    "ngrok": {
+        "oauth_provider": "google",
+        "oauth_allow_domains": "asdf.com",
+        "oauth_allow_emails": "asdf@asdf.com"
+    }
+}
+```
+
+To add an authtoken instead of using the NGROK_AUTHTOKEN environment variable:
+```json
+{
+    "ngrok": {
+        "authtoken": "<token>",
+        "authtoken_from_env":false
+    }
+}
+```

extensions/ngrok/requirements.txt

@@ -0,0 +1 @@
+ngrok==0.*

extensions/ngrok/script.py

@@ -0,0 +1,36 @@
+# Adds ngrok ingress, to use add `--extension ngrok` to the command line options
+#
+# Parameters can be customized in settings.json of webui, e.g.:
+# {"ngrok": {"basic_auth":"user:password"} }
+# or
+# {"ngrok": {"oauth_provider":"google", "oauth_allow_emails":["asdf@asdf.com"]} }
+#
+# See this example for full list of options: https://github.com/ngrok/ngrok-py/blob/main/examples/ngrok-connect-full.py
+# or the README.md in this directory.
+
+import logging
+from modules import shared
+
+# Pick up host/port command line arguments
+host = shared.args.listen_host if shared.args.listen_host and shared.args.listen else '127.0.0.1'
+port = shared.args.listen_port if shared.args.listen_port else '7860'
+
+# Default options
+options = {
+    'addr': f"{host}:{port}",
+    'authtoken_from_env': True,
+    'session_metadata': 'text-generation-webui',
+}
+
+
+def ui():
+    settings = shared.settings.get("ngrok")
+    if settings:
+        options.update(settings)
+
+    try:
+        import ngrok
+        tunnel = ngrok.connect(**options)
+        logging.info(f"Ingress established at: {tunnel.url()}")
+    except ModuleNotFoundError:
+        logging.error("===> ngrok library not found, please run `pip install -r extensions/ngrok/requirements.txt`")

extensions/openai/cache_embedding_model.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+# preload the embedding model, useful for Docker images to prevent re-download on config change
+# Dockerfile:
+# ENV OPENEDAI_EMBEDDING_MODEL="sentence-transformers/all-mpnet-base-v2"  # Optional
+# RUN python3 cache_embedded_model.py
+import os
+
+import sentence_transformers
+
+st_model = os.environ.get("OPENEDAI_EMBEDDING_MODEL", "sentence-transformers/all-mpnet-base-v2")
+model = sentence_transformers.SentenceTransformer(st_model)

extensions/openai/completions.py

@@ -0,0 +1,532 @@
+import base64
+import copy
+import re
+import time
+from collections import deque
+from io import BytesIO
+
+import requests
+import tiktoken
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from transformers import LogitsProcessor, LogitsProcessorList
+
+from extensions.openai.errors import InvalidRequestError
+from extensions.openai.utils import debug_msg
+from modules import shared
+from modules.chat import (
+    generate_chat_prompt,
+    generate_chat_reply,
+    load_character_memoized,
+    load_instruction_template_memoized
+)
+from modules.presets import load_preset_memoized
+from modules.text_generation import decode, encode, generate_reply
+
+
+class LogitsBiasProcessor(LogitsProcessor):
+    def __init__(self, logit_bias={}):
+        self.logit_bias = logit_bias
+        if self.logit_bias:
+            self.keys = list([int(key) for key in self.logit_bias.keys()])
+            values = [self.logit_bias[str(key)] for key in self.keys]
+            self.values = torch.tensor(values, dtype=torch.float, device=shared.model.device)
+            debug_msg(f"{self})")
+
+    def __call__(self, input_ids: torch.LongTensor, logits: torch.FloatTensor) -> torch.FloatTensor:
+        if self.logit_bias:
+            debug_msg(logits[0, self.keys], " + ", self.values)
+            logits[0, self.keys] += self.values
+            debug_msg(" --> ", logits[0, self.keys])
+            debug_msg(" max/min ", float(torch.max(logits[0])), float(torch.min(logits[0])))
+
+        return logits
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__}(logit_bias={self.logit_bias})>"
+
+
+class LogprobProcessor(LogitsProcessor):
+    def __init__(self, logprobs=None):
+        self.logprobs = logprobs
+        self.token_alternatives = {}
+
+    def __call__(self, input_ids: torch.LongTensor, logits: torch.FloatTensor) -> torch.FloatTensor:
+        if self.logprobs is not None:  # 0-5
+            log_e_probabilities = F.log_softmax(logits, dim=1)
+            top_values, top_indices = torch.topk(log_e_probabilities, k=self.logprobs + 1)
+            top_tokens = [decode(tok) for tok in top_indices[0]]
+            top_probs = [float(x) for x in top_values[0]]
+            self.token_alternatives = dict(zip(top_tokens, top_probs))
+            debug_msg(repr(self))
+
+        return logits
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__}(logprobs={self.logprobs}, token_alternatives={self.token_alternatives})>"
+
+
+def convert_logprobs_to_tiktoken(model, logprobs):
+    # more problems than it's worth.
+    # try:
+    #     encoder = tiktoken.encoding_for_model(model)
+    #     # just pick the first one if it encodes to multiple tokens... 99.9% not required and maybe worse overall.
+    #     return dict([(encoder.decode([encoder.encode(token)[0]]), prob) for token, prob in logprobs.items()])
+    # except KeyError:
+    #     # assume native tokens if we can't find the tokenizer
+    #     return logprobs
+
+    return logprobs
+
+
+def process_parameters(body, is_legacy=False):
+    generate_params = body
+    max_tokens_str = 'length' if is_legacy else 'max_tokens'
+    generate_params['max_new_tokens'] = body.pop(max_tokens_str)
+    if generate_params['truncation_length'] == 0:
+        generate_params['truncation_length'] = shared.settings['truncation_length']
+
+    if body['preset'] is not None:
+        preset = load_preset_memoized(body['preset'])
+        generate_params.update(preset)
+
+    generate_params['custom_stopping_strings'] = []
+    if 'stop' in body:  # str or array, max len 4 (ignored)
+        if isinstance(body['stop'], str):
+            generate_params['custom_stopping_strings'] = [body['stop']]
+        elif isinstance(body['stop'], list):
+            generate_params['custom_stopping_strings'] = body['stop']
+
+    logits_processor = []
+    logit_bias = body.get('logit_bias', None)
+    if logit_bias:  # {str: float, ...}
+        # XXX convert tokens from tiktoken based on requested model
+        # Ex.: 'logit_bias': {'1129': 100, '11442': 100, '16243': 100}
+        try:
+            encoder = tiktoken.encoding_for_model(generate_params['model'])
+            new_logit_bias = {}
+            for logit, bias in logit_bias.items():
+                for x in encode(encoder.decode([int(logit)]), add_special_tokens=False)[0]:
+                    if int(x) in [0, 1, 2, 29871]:  # XXX LLAMA tokens
+                        continue
+
+                    new_logit_bias[str(int(x))] = bias
+            debug_msg('logit_bias_map', logit_bias, '->', new_logit_bias)
+            logit_bias = new_logit_bias
+        except KeyError:
+            pass  # assume native tokens if we can't find the tokenizer
+
+        logits_processor = [LogitsBiasProcessor(logit_bias)]
+
+    logprobs = None  # coming to chat eventually
+    if 'logprobs' in body:
+        logprobs = body.get('logprobs', 0)  # maybe cap at topk? don't clamp 0-5.
+        generate_params['logprob_proc'] = LogprobProcessor(logprobs)
+        logits_processor.extend([generate_params['logprob_proc']])
+    else:
+        logprobs = None
+
+    if logits_processor:  # requires logits_processor support
+        generate_params['logits_processor'] = LogitsProcessorList(logits_processor)
+
+    return generate_params
+
+
+def convert_history(history):
+    '''
+    Chat histories in this program are in the format [message, reply].
+    This function converts OpenAI histories to that format.
+    '''
+    chat_dialogue = []
+    current_message = ""
+    current_reply = ""
+    user_input = ""
+    system_message = ""
+
+    for entry in history:
+        if "image_url" in entry:
+            image_url = entry['image_url']
+            if "base64" in image_url:
+                image_url = re.sub('^data:image/.+;base64,', '', image_url)
+                img = Image.open(BytesIO(base64.b64decode(image_url)))
+            else:
+                try:
+                    my_res = requests.get(image_url)
+                    img = Image.open(BytesIO(my_res.content))
+                except Exception:
+                    raise 'Image cannot be loaded from the URL!'
+
+            buffered = BytesIO()
+            img.save(buffered, format="JPEG")
+            img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+            content = f'<img src="data:image/jpeg;base64,{img_str}">'
+        else:
+            content = entry["content"]
+
+        role = entry["role"]
+
+        if role == "user":
+            user_input = content
+            if current_message:
+                chat_dialogue.append([current_message, ''])
+                current_message = ""
+            current_message = content
+        elif role == "assistant":
+            current_reply = content
+            if current_message:
+                chat_dialogue.append([current_message, current_reply])
+                current_message = ""
+                current_reply = ""
+            else:
+                chat_dialogue.append(['', current_reply])
+        elif role == "system":
+            system_message = content
+
+    # if current_message:
+    #     chat_dialogue.append([current_message, ''])
+
+    return user_input, system_message, {'internal': chat_dialogue, 'visible': copy.deepcopy(chat_dialogue)}
+
+
+def chat_completions_common(body: dict, is_legacy: bool = False, stream=False) -> dict:
+    if body.get('functions', []):
+        raise InvalidRequestError(message="functions is not supported.", param='functions')
+
+    if body.get('function_call', ''):
+        raise InvalidRequestError(message="function_call is not supported.", param='function_call')
+
+    if 'messages' not in body:
+        raise InvalidRequestError(message="messages is required", param='messages')
+
+    messages = body['messages']
+    for m in messages:
+        if 'role' not in m:
+            raise InvalidRequestError(message="messages: missing role", param='messages')
+        elif m['role'] == 'function':
+            raise InvalidRequestError(message="role: function is not supported.", param='messages')
+
+        if 'content' not in m and "image_url" not in m:
+            raise InvalidRequestError(message="messages: missing content", param='messages')
+
+    # Chat Completions
+    object_type = 'chat.completions' if not stream else 'chat.completions.chunk'
+    created_time = int(time.time())
+    cmpl_id = "chatcmpl-%d" % (int(time.time() * 1000000000))
+    resp_list = 'data' if is_legacy else 'choices'
+
+    # generation parameters
+    generate_params = process_parameters(body, is_legacy=is_legacy)
+    continue_ = body['continue_']
+
+    # Instruction template
+    if body['instruction_template_str']:
+        instruction_template_str = body['instruction_template_str']
+    elif body['instruction_template']:
+        instruction_template = body['instruction_template']
+        instruction_template = "Alpaca" if instruction_template == "None" else instruction_template
+        instruction_template_str = load_instruction_template_memoized(instruction_template)
+    else:
+        instruction_template_str = shared.settings['instruction_template_str']
+
+    chat_template_str = body['chat_template_str'] or shared.settings['chat_template_str']
+    chat_instruct_command = body['chat_instruct_command'] or shared.settings['chat-instruct_command']
+
+    # Chat character
+    character = body['character'] or shared.settings['character']
+    character = "Assistant" if character == "None" else character
+    name1 = body['user_name'] or shared.settings['name1']
+    name1, name2, _, greeting, context = load_character_memoized(character, name1, '')
+    name2 = body['bot_name'] or name2
+    context = body['context'] or context
+    greeting = body['greeting'] or greeting
+
+    # History
+    user_input, custom_system_message, history = convert_history(messages)
+
+    generate_params.update({
+        'mode': body['mode'],
+        'name1': name1,
+        'name2': name2,
+        'context': context,
+        'greeting': greeting,
+        'instruction_template_str': instruction_template_str,
+        'custom_system_message': custom_system_message,
+        'chat_template_str': chat_template_str,
+        'chat-instruct_command': chat_instruct_command,
+        'history': history,
+        'stream': stream
+    })
+
+    max_tokens = generate_params['max_new_tokens']
+    if max_tokens in [None, 0]:
+        generate_params['max_new_tokens'] = 512
+        generate_params['auto_max_new_tokens'] = True
+
+    requested_model = generate_params.pop('model')
+    logprob_proc = generate_params.pop('logprob_proc', None)
+
+    def chat_streaming_chunk(content):
+        # begin streaming
+        chunk = {
+            "id": cmpl_id,
+            "object": object_type,
+            "created": created_time,
+            "model": shared.model_name,
+            resp_list: [{
+                "index": 0,
+                "finish_reason": None,
+                # So yeah... do both methods? delta and messages.
+                "message": {'role': 'assistant', 'content': content},
+                "delta": {'role': 'assistant', 'content': content},
+            }],
+        }
+
+        if logprob_proc:  # not official for chat yet
+            top_logprobs = convert_logprobs_to_tiktoken(model=requested_model, logprobs=logprob_proc.token_alternatives)
+            chunk[resp_list][0]["logprobs"] = {'top_logprobs': [top_logprobs]}
+        # else:
+        #    chunk[resp_list][0]["logprobs"] = None
+        return chunk
+
+    if stream:
+        yield chat_streaming_chunk('')
+
+    # generate reply #######################################
+    prompt = generate_chat_prompt(user_input, generate_params)
+    token_count = len(encode(prompt)[0])
+    debug_msg({'prompt': prompt, 'generate_params': generate_params})
+
+    generator = generate_chat_reply(
+        user_input, generate_params, regenerate=False, _continue=continue_, loading_message=False)
+
+    answer = ''
+    seen_content = ''
+    completion_token_count = 0
+
+    for a in generator:
+        answer = a['internal'][-1][1]
+        if stream:
+            len_seen = len(seen_content)
+            new_content = answer[len_seen:]
+
+            if not new_content or chr(0xfffd) in new_content:  # partial unicode character, don't send it yet.
+                continue
+
+            seen_content = answer
+            chunk = chat_streaming_chunk(new_content)
+            yield chunk
+
+    completion_token_count = len(encode(answer)[0])
+    stop_reason = "stop"
+    if token_count + completion_token_count >= generate_params['truncation_length'] or completion_token_count >= generate_params['max_new_tokens']:
+        stop_reason = "length"
+
+    if stream:
+        chunk = chat_streaming_chunk('')
+        chunk[resp_list][0]['finish_reason'] = stop_reason
+        chunk['usage'] = {
+            "prompt_tokens": token_count,
+            "completion_tokens": completion_token_count,
+            "total_tokens": token_count + completion_token_count
+        }
+
+        yield chunk
+    else:
+        resp = {
+            "id": cmpl_id,
+            "object": object_type,
+            "created": created_time,
+            "model": shared.model_name,
+            resp_list: [{
+                "index": 0,
+                "finish_reason": stop_reason,
+                "message": {"role": "assistant", "content": answer}
+            }],
+            "usage": {
+                "prompt_tokens": token_count,
+                "completion_tokens": completion_token_count,
+                "total_tokens": token_count + completion_token_count
+            }
+        }
+        if logprob_proc:  # not official for chat yet
+            top_logprobs = convert_logprobs_to_tiktoken(model=requested_model, logprobs=logprob_proc.token_alternatives)
+            resp[resp_list][0]["logprobs"] = {'top_logprobs': [top_logprobs]}
+        # else:
+        #     resp[resp_list][0]["logprobs"] = None
+
+        yield resp
+
+
+def completions_common(body: dict, is_legacy: bool = False, stream=False):
+    object_type = 'text_completion.chunk' if stream else 'text_completion'
+    created_time = int(time.time())
+    cmpl_id = "conv-%d" % (int(time.time() * 1000000000))
+    resp_list = 'data' if is_legacy else 'choices'
+
+    prompt_str = 'context' if is_legacy else 'prompt'
+
+    # ... encoded as a string, array of strings, array of tokens, or array of token arrays.
+    if prompt_str not in body:
+        raise InvalidRequestError("Missing required input", param=prompt_str)
+
+    # common params
+    generate_params = process_parameters(body, is_legacy=is_legacy)
+    max_tokens = generate_params['max_new_tokens']
+    generate_params['stream'] = stream
+    requested_model = generate_params.pop('model')
+    logprob_proc = generate_params.pop('logprob_proc', None)
+    suffix = body['suffix'] if body['suffix'] else ''
+    echo = body['echo']
+
+    if not stream:
+        prompt_arg = body[prompt_str]
+        if isinstance(prompt_arg, str) or (isinstance(prompt_arg, list) and isinstance(prompt_arg[0], int)):
+            prompt_arg = [prompt_arg]
+
+        resp_list_data = []
+        total_completion_token_count = 0
+        total_prompt_token_count = 0
+
+        for idx, prompt in enumerate(prompt_arg, start=0):
+            if isinstance(prompt[0], int):
+                # token lists
+                if requested_model == shared.model_name:
+                    prompt = decode(prompt)[0]
+                else:
+                    try:
+                        encoder = tiktoken.encoding_for_model(requested_model)
+                        prompt = encoder.decode(prompt)
+                    except KeyError:
+                        prompt = decode(prompt)[0]
+
+            prefix = prompt if echo else ''
+            token_count = len(encode(prompt)[0])
+            total_prompt_token_count += token_count
+
+            # generate reply #######################################
+            debug_msg({'prompt': prompt, 'generate_params': generate_params})
+            generator = generate_reply(prompt, generate_params, is_chat=False)
+            answer = ''
+
+            for a in generator:
+                answer = a
+
+            completion_token_count = len(encode(answer)[0])
+            total_completion_token_count += completion_token_count
+            stop_reason = "stop"
+            if token_count + completion_token_count >= generate_params['truncation_length'] or completion_token_count >= max_tokens:
+                stop_reason = "length"
+
+            respi = {
+                "index": idx,
+                "finish_reason": stop_reason,
+                "text": prefix + answer + suffix,
+                "logprobs": {'top_logprobs': [logprob_proc.token_alternatives]} if logprob_proc else None,
+            }
+
+            resp_list_data.extend([respi])
+
+        resp = {
+            "id": cmpl_id,
+            "object": object_type,
+            "created": created_time,
+            "model": shared.model_name,
+            resp_list: resp_list_data,
+            "usage": {
+                "prompt_tokens": total_prompt_token_count,
+                "completion_tokens": total_completion_token_count,
+                "total_tokens": total_prompt_token_count + total_completion_token_count
+            }
+        }
+
+        yield resp
+    else:
+        prompt = body[prompt_str]
+        if isinstance(prompt, list):
+            if prompt and isinstance(prompt[0], int):
+                try:
+                    encoder = tiktoken.encoding_for_model(requested_model)
+                    prompt = encoder.decode(prompt)
+                except KeyError:
+                    prompt = decode(prompt)[0]
+            else:
+                raise InvalidRequestError(message="API Batched generation not yet supported.", param=prompt_str)
+
+        prefix = prompt if echo else ''
+        token_count = len(encode(prompt)[0])
+
+        def text_streaming_chunk(content):
+            # begin streaming
+            chunk = {
+                "id": cmpl_id,
+                "object": object_type,
+                "created": created_time,
+                "model": shared.model_name,
+                resp_list: [{
+                    "index": 0,
+                    "finish_reason": None,
+                    "text": content,
+                    "logprobs": {'top_logprobs': [logprob_proc.token_alternatives]} if logprob_proc else None,
+                }],
+            }
+
+            return chunk
+
+        yield text_streaming_chunk(prefix)
+
+        # generate reply #######################################
+        debug_msg({'prompt': prompt, 'generate_params': generate_params})
+        generator = generate_reply(prompt, generate_params, is_chat=False)
+
+        answer = ''
+        seen_content = ''
+        completion_token_count = 0
+
+        for a in generator:
+            answer = a
+
+            len_seen = len(seen_content)
+            new_content = answer[len_seen:]
+
+            if not new_content or chr(0xfffd) in new_content:  # partial unicode character, don't send it yet.
+                continue
+
+            seen_content = answer
+            chunk = text_streaming_chunk(new_content)
+            yield chunk
+
+        completion_token_count = len(encode(answer)[0])
+        stop_reason = "stop"
+        if token_count + completion_token_count >= generate_params['truncation_length'] or completion_token_count >= max_tokens:
+            stop_reason = "length"
+
+        chunk = text_streaming_chunk(suffix)
+        chunk[resp_list][0]["finish_reason"] = stop_reason
+        chunk["usage"] = {
+            "prompt_tokens": token_count,
+            "completion_tokens": completion_token_count,
+            "total_tokens": token_count + completion_token_count
+        }
+
+        yield chunk
+
+
+def chat_completions(body: dict, is_legacy: bool = False) -> dict:
+    generator = chat_completions_common(body, is_legacy, stream=False)
+    return deque(generator, maxlen=1).pop()
+
+
+def stream_chat_completions(body: dict, is_legacy: bool = False):
+    for resp in chat_completions_common(body, is_legacy, stream=True):
+        yield resp
+
+
+def completions(body: dict, is_legacy: bool = False) -> dict:
+    generator = completions_common(body, is_legacy, stream=False)
+    return deque(generator, maxlen=1).pop()
+
+
+def stream_completions(body: dict, is_legacy: bool = False):
+    for resp in completions_common(body, is_legacy, stream=True):
+        yield resp

extensions/openai/embeddings.py

@@ -0,0 +1,98 @@
+import os
+
+import numpy as np
+from transformers import AutoModel
+
+from extensions.openai.errors import ServiceUnavailableError
+from extensions.openai.utils import debug_msg, float_list_to_base64
+from modules.logging_colors import logger
+
+embeddings_params_initialized = False
+
+
+def initialize_embedding_params():
+    '''
+    using 'lazy loading' to avoid circular import
+    so this function will be executed only once
+    '''
+    global embeddings_params_initialized
+    if not embeddings_params_initialized:
+        from extensions.openai.script import params
+
+        global st_model, embeddings_model, embeddings_device
+
+        st_model = os.environ.get("OPENEDAI_EMBEDDING_MODEL", params.get('embedding_model', 'all-mpnet-base-v2'))
+        embeddings_model = None
+        # OPENEDAI_EMBEDDING_DEVICE: auto (best or cpu), cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, ideep, hip, ve, fpga, ort, xla, lazy, vulkan, mps, meta, hpu, mtia, privateuseone
+        embeddings_device = os.environ.get("OPENEDAI_EMBEDDING_DEVICE", params.get('embedding_device', 'cpu'))
+        if embeddings_device.lower() == 'auto':
+            embeddings_device = None
+
+        embeddings_params_initialized = True
+
+
+def load_embedding_model(model: str):
+    try:
+        from sentence_transformers import SentenceTransformer
+    except ModuleNotFoundError:
+        logger.error("The sentence_transformers module has not been found. Please install it manually with pip install -U sentence-transformers.")
+        raise ModuleNotFoundError
+
+    initialize_embedding_params()
+    global embeddings_device, embeddings_model
+    try:
+        print(f"Try embedding model: {model} on {embeddings_device}")
+        if 'jina-embeddings' in model:
+            embeddings_model = AutoModel.from_pretrained(model, trust_remote_code=True)  # trust_remote_code is needed to use the encode method
+            embeddings_model = embeddings_model.to(embeddings_device)
+        else:
+            embeddings_model = SentenceTransformer(model, device=embeddings_device)
+
+        print(f"Loaded embedding model: {model}")
+    except Exception as e:
+        embeddings_model = None
+        raise ServiceUnavailableError(f"Error: Failed to load embedding model: {model}", internal_message=repr(e))
+
+
+def get_embeddings_model():
+    initialize_embedding_params()
+    global embeddings_model, st_model
+    if st_model and not embeddings_model:
+        load_embedding_model(st_model)  # lazy load the model
+
+    return embeddings_model
+
+
+def get_embeddings_model_name() -> str:
+    initialize_embedding_params()
+    global st_model
+    return st_model
+
+
+def get_embeddings(input: list) -> np.ndarray:
+    model = get_embeddings_model()
+    debug_msg(f"embedding model : {model}")
+    embedding = model.encode(input, convert_to_numpy=True, normalize_embeddings=True, convert_to_tensor=False)
+    debug_msg(f"embedding result : {embedding}")  # might be too long even for debug, use at you own will
+    return embedding
+
+
+def embeddings(input: list, encoding_format: str) -> dict:
+    embeddings = get_embeddings(input)
+    if encoding_format == "base64":
+        data = [{"object": "embedding", "embedding": float_list_to_base64(emb), "index": n} for n, emb in enumerate(embeddings)]
+    else:
+        data = [{"object": "embedding", "embedding": emb.tolist(), "index": n} for n, emb in enumerate(embeddings)]
+
+    response = {
+        "object": "list",
+        "data": data,
+        "model": st_model,  # return the real model
+        "usage": {
+            "prompt_tokens": 0,
+            "total_tokens": 0,
+        }
+    }
+
+    debug_msg(f"Embeddings return size: {len(embeddings[0])}, number: {len(embeddings)}")
+    return response

extensions/openai/errors.py

@@ -0,0 +1,31 @@
+class OpenAIError(Exception):
+    def __init__(self, message=None, code=500, internal_message=''):
+        self.message = message
+        self.code = code
+        self.internal_message = internal_message
+
+    def __repr__(self):
+        return "%s(message=%r, code=%d)" % (
+            self.__class__.__name__,
+            self.message,
+            self.code,
+        )
+
+
+class InvalidRequestError(OpenAIError):
+    def __init__(self, message, param, code=400, internal_message=''):
+        super().__init__(message, code, internal_message)
+        self.param = param
+
+    def __repr__(self):
+        return "%s(message=%r, code=%d, param=%s)" % (
+            self.__class__.__name__,
+            self.message,
+            self.code,
+            self.param,
+        )
+
+
+class ServiceUnavailableError(OpenAIError):
+    def __init__(self, message="Service unavailable, please try again later.", code=503, internal_message=''):
+        super().__init__(message, code, internal_message)