File size: 40,162 Bytes
a509ff9
 
90290aa
a509ff9
83610fc
a509ff9
 
a33aa9b
0e02b5f
 
 
 
 
f1da0db
90290aa
5d5efcb
 
 
 
 
9c6cacb
 
83610fc
b4bf2a9
8b24305
a077145
8b24305
 
a33aa9b
a509ff9
d96b2be
611defa
d96b2be
611defa
 
71058a8
 
 
 
83610fc
5232b3e
71058a8
 
d96b2be
 
 
 
 
611defa
 
71058a8
d96b2be
611defa
 
 
 
 
d96b2be
5232b3e
d96b2be
 
 
 
71058a8
 
 
 
d96b2be
 
 
 
a41a44d
d96b2be
 
 
 
 
611defa
 
 
72cd75e
 
d96b2be
611defa
 
3f3bfef
611defa
 
 
3f3bfef
 
0e986bc
 
 
3f3bfef
0e986bc
 
c8c8a26
 
 
 
 
 
 
 
 
 
 
 
 
9c6cacb
c8c8a26
9c6cacb
c8c8a26
9c6cacb
 
d96b2be
 
 
72cd75e
 
83610fc
 
 
 
 
 
f1da0db
a077145
 
 
 
 
 
 
 
 
 
f1da0db
a077145
 
 
 
 
e7aeeed
a077145
 
e7aeeed
a41a44d
f1da0db
a077145
 
 
 
e7aeeed
f1da0db
22f7149
8b24305
22f7149
 
 
 
 
 
b4bf2a9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f57539d
a077145
a509ff9
1e4bd6c
 
a509ff9
90290aa
a509ff9
 
 
1e4bd6c
 
 
a509ff9
90290aa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0e02b5f
72cd75e
 
 
 
 
 
 
 
 
 
 
 
 
b4bf2a9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9c6cacb
 
 
 
 
 
 
 
 
 
 
 
5d5efcb
 
9c6cacb
5d5efcb
 
 
 
 
9c6cacb
5d5efcb
 
0e02b5f
 
 
90290aa
0e02b5f
 
 
 
 
 
 
 
 
 
 
 
 
 
71058a8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0e02b5f
 
90290aa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
71058a8
83610fc
71058a8
83610fc
71058a8
83610fc
71058a8
c03c42a
 
71058a8
 
 
 
 
 
c03c42a
71058a8
 
42702f4
71058a8
 
 
 
 
cb15d1f
 
 
 
83610fc
cb15d1f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
83610fc
 
 
 
 
 
 
 
 
 
 
 
b4bf2a9
83610fc
 
 
 
 
 
 
 
 
 
 
b4bf2a9
83610fc
b4bf2a9
 
 
 
 
 
 
 
 
 
 
83610fc
b4bf2a9
 
 
 
 
 
 
 
83610fc
 
b4bf2a9
83610fc
b4bf2a9
 
 
83610fc
b4bf2a9
 
 
 
83610fc
b4bf2a9
83610fc
 
f1da0db
cb15d1f
83610fc
f1da0db
 
 
cb15d1f
f1da0db
cb15d1f
83610fc
0e02b5f
71058a8
0e02b5f
 
 
 
 
 
 
 
71058a8
 
 
 
0e02b5f
e7aeeed
 
90290aa
71058a8
83610fc
cb15d1f
 
e7aeeed
d3650cd
e7aeeed
 
 
 
 
f1da0db
e7aeeed
 
d3650cd
90290aa
 
d3650cd
 
 
 
 
f1da0db
 
 
d3650cd
e7aeeed
90290aa
3f3bfef
a077145
 
e7aeeed
a077145
e7aeeed
a077145
 
0e02b5f
 
 
 
 
 
 
 
 
 
 
 
075d09e
72cd75e
 
 
 
90290aa
 
 
e7aeeed
90290aa
 
71058a8
90290aa
 
e7aeeed
 
f1da0db
e7aeeed
cb15d1f
71058a8
cb15d1f
f1da0db
0e02b5f
e7aeeed
 
 
 
 
 
f1da0db
e7aeeed
 
5232b3e
b4bf2a9
e7aeeed
 
 
 
 
 
 
5232b3e
f1da0db
b4bf2a9
f1da0db
e7aeeed
 
 
f1da0db
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5b174ea
 
f1da0db
5b174ea
f1da0db
 
 
9c6cacb
e7aeeed
611defa
e7aeeed
22f7149
 
611defa
 
f1da0db
611defa
 
f1da0db
e7aeeed
3f3bfef
611defa
e7aeeed
83610fc
 
f1da0db
 
e7aeeed
90290aa
83610fc
 
 
 
f57539d
b4bf2a9
 
 
 
 
f57539d
 
 
 
 
 
83610fc
 
 
f57539d
 
 
 
 
 
 
 
 
cb15d1f
611defa
90290aa
9c6cacb
 
 
 
 
 
5d5efcb
83610fc
611defa
83610fc
5d5efcb
 
83610fc
5d5efcb
f4478ff
d3650cd
 
f1da0db
d3650cd
65c472b
8aafa85
f57539d
dcaa01a
b4bf2a9
8b24305
 
 
62e022b
 
f57539d
dcaa01a
f4478ff
d3650cd
 
 
 
 
83610fc
 
 
 
 
cb15d1f
d3650cd
 
 
 
9c6cacb
 
 
 
79c8c7e
 
9c6cacb
 
5d5efcb
 
 
83610fc
5d5efcb
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
import gradio as gr
from sklearn.metrics.pairwise import cosine_similarity
from scipy.sparse import csr_matrix
import numpy as np
import joblib
from joblib import load
import h5py
from io import BytesIO
import csv
import re
import random
import compress_fasttext
from collections import OrderedDict
from lark import Lark, Tree, Token
from lark.exceptions import ParseError
import json
import zipfile
from PIL import Image
import io
import os
import glob
import itertools
from itertools import islice
from pathlib import Path
import logging

# Set up logging
logging.basicConfig(filename='error.log', level=logging.DEBUG, format='%(asctime)s %(levelname)s:%(message)s')


faq_content="""
# Questions:

## What is the purpose of this tool?

Since Stable Diffusion's initial release in 2022, users have developed a myriad of fine-tuned text to image models, each with unique "linguistic" preferences depending on the data from which it was fine-tuned.
Some models react best when prompted with verbose scene descriptions akin to DALL-E, while others fine-tuned on images scraped from popular image boards understand those boards' tag sets.
This tool serves as a linguistic bridge to the e621 image board tag lexicon, on which many popular models such as Fluffyrock, Fluffusion, and Pony Diffusion v6 were trained.

When you enter a txt2img prompt and press the "submit" button, Prompt Squirrel parses your prompt and checks that all your tags are valid e621 tags.
If it finds any that are not, it recommends some valid e621 tags you can use to replace them in the "Unknown Tags" section.
Additionally, in the "Top Artists" text box, it lists the artists who would most likely draw an image having the set of tags you provided.
This is useful to align your prompt with the expected input to an e621-trained model.

## Does input order matter?

No

## Should I use underscores or spaces in the input tags?

As a rule, e621-trained models replace underscores in tags with spaces, so spaces are preferred.

## Can I use parentheses or weights as in the Stable Diffusion Automatic1111 WebUI?

Yes, but only '(' and ')' and numerical weights, and all of these things are ignored in all calculations.  The main benefit of this is that you can copy/paste prompts from one program to another with minimal editing.  
An example that illustrates acceptable parentheses and weight formatting is:
((sunset over the mountains)), (clear sky:1.5), ((eagle flying high:2.0)), river, (fish swimming in the river:1.2), (campfire, (marshmallows:2.1):1.3), stars in the sky, ((full moon:1.8)), (wolf howling:1.7)

## Why are some valid tags marked as "unknown", and why don't some artists ever get returned?

Some data is excluded from consideration if it did not occur frequently enough in the sample from which the application makes its calculations.
If an artist or tag is too infrequent, we might not think we have enough data to make predictions about it.

## Why do some suggested tags not have summaries or wiki links?

Both of these features are extracted from the tag wiki pages, but some valid e621 tags do not have wiki pages.

## Are there any special tags?

Yes.  We normalized the favorite counts of each image to a range of 0-9, with 0 being the lowest favcount, and 9 being the highest.
You can include any of these special tags: "score:0", "score:1", "score:2", "score:3", "score:4", "score:5", "score:6", "score:7", "score:8", "score:9"
in your list to bias the output toward artists with higher or lower scoring images.

## Are there any other special tricks?

Yes.  If you want to more strongly bias the artist output toward a specific tag, you can just list it multiple times.  
So for example, the query "red fox, red fox, red fox, score:7" will yield a list of artists who are more strongly associated with the tag "red fox"
than the query "red fox, score:7".  

## Why is this space tagged "not-for-all-audience"
The "not-for-all-audience" tag informs users that this tool's text output is derived from e621.net data for tag prediction and completion.
The app will try not to display nsfw tags unless the "Allow NSFW Tags" is checked, but the filter is not perfect.

## How does the tag corrector work?

We collect the tag sets from over 4 million e621 posts, treating the tag set from each image as an individual document.
We then randomly replace about 10% of the tags in each document with a randomly selected alias from e621's list of aliases for the tag 
(e.g. "canine" gets replaced with one of {k9,canines,mongrel,cannine,cnaine,feral_canine,anthro_canine}).
We then train a FastText (https://fasttext.cc/) model on the documents.  The result of this training is a function that maps arbitrary words to vectors such that
the vector for a tag and the vectors for its aliases are all close together (because the model has seen them in similar contexts).
Since the lists of aliases contain misspellings and rephrasings of tags, the model should be robust to these kinds of problems as long as they are not too dissimilar from the alias lists.  

To enhance the tag corrector further, we leverage conditional probabilities to refine our predictions. 
Using the same 4 million post dataset, we calculate the conditional probability of each tag given the context of other tags appearing within the same document.
This is done by creating a co-occurrence matrix from our dataset, which records how frequently each pair of tags appears together across all documents.
By considering the context in which tags are used, we can now not only correct misspellings and rephrasings but also make more contextually relevant suggestions.
The "similarity weight" slider controls how much weight these conditional probabilities are given vs how much weight the FastText similarity model is given when suggesting replacements for invalid tags.
A similarity weight slider value of 0 means that only the FastText model's predictions will be used to calculate similarity scores, and a value of 1 means only the conditional probabilities are used (although the FastText model is still used to trim the list of candidates).

## How is the artist list calculated?

Each artist is represented by a "pseudo-document" composed of all the tags from their uploaded images, treating these tags similarly to words in a text document. 
Similarly, when you input a set of tags, the system creates a pseudo-document for your query out of all the tags. 
It then uses a technique called cosine similarity to compare your tags against each artist's collection, essentially finding which artist's tags are most "similar" to yours.
This method helps identify artists whose work is closely aligned with the themes or elements you're interested in.
For those curious about the underlying mechanics of comparing text-like data, we employ the TF-IDF (Term Frequency-Inverse Document Frequency) method, a standard approach in information retrieval. 
You can read more about TF-IDF on its [Wikipedia page](https://en.wikipedia.org/wiki/Tf%E2%80%93idf).  

## How do the sample images work?

In the first row of galleries, for each artist in the dataset, we generated a sample image with the model Fluffyrock Unleashed using the prompt "by artist, soyjak, anthro, male, bust portrait, meme, grin" where "artist" is the name of an artist.
The simplicity of the prompt, the the simplicty of the default style, and the recognizability of the character make it easier to understand how artist names affect generated image styles.
The image on the left captioned "No Artist" was generated with the same prompt, but with no artist name.
You should compare all the images to the first to see how the artist names affect the output.
Each subsequent row of images was generated using the same process, but with a different prompt.
See SamplePrompts.csv for the list of prompts used and their descriptions.
"""


nsfw_threshold = 0.95  # Assuming the threshold value is defined here

css = """
.scrollable-content {
    max-height: 500px;
    overflow-y: auto;
}
"""

grammar=r"""
!start: (prompt | /[][():]/+)*
prompt: (emphasized | plain | comma | WHITESPACE)*
!emphasized: "(" prompt ")"
        | "(" prompt ":" [WHITESPACE] NUMBER [WHITESPACE] ")"
comma: ","
WHITESPACE: /\s+/
plain: /([^,\\\[\]():|]|\\.)+/
%import common.SIGNED_NUMBER -> NUMBER
"""

# Initialize the parser
parser = Lark(grammar, start='start')

# Function to extract tags
def extract_tags(tree):
    tags_with_positions = []
    def _traverse(node):
        if isinstance(node, Token) and node.type == '__ANON_1':
            tag_position = node.start_pos
            tag_text = node.value
            tags_with_positions.append((tag_text, tag_position, "tag"))
        elif not isinstance(node, Token):
            for child in node.children:
                _traverse(child)
    _traverse(tree)
    return tags_with_positions
    

special_tags = ["score:0", "score:1", "score:2", "score:3", "score:4", "score:5", "score:6", "score:7", "score:8", "score:9", "rating:s", "rating:q", "rating:e"]
def remove_special_tags(original_string):
    tags = [tag.strip() for tag in original_string.split(",")]
    remaining_tags = [tag for tag in tags if tag not in special_tags]
    removed_tags = [tag for tag in tags if tag in special_tags]
    return ", ".join(remaining_tags), removed_tags
    

# Define a function to load all necessary components
def load_model_components(file_path):
    # Ensure the file path is a Path object for robust path handling
    file_path = Path(file_path)
    
    # Check if the file exists
    if not file_path.is_file():
        raise FileNotFoundError(f"The specified joblib file was not found: {file_path}")
    
    # Load all the model components from the joblib file
    model_components = joblib.load(file_path)
    
    # Create a reverse mapping from row index to tag
    if 'tag_to_row_index' in model_components:
        model_components['row_to_tag'] = {idx: tag for tag, idx in model_components['tag_to_row_index'].items()}
    
    return model_components
# Load all components at the start
tf_idf_components = load_model_components('tf_idf_files_420.joblib')
    
# Load the model and data once at startup
with h5py.File('complete_artist_data.hdf5', 'r') as f:
    # Deserialize the vectorizer
    vectorizer_bytes = f['vectorizer'][()].tobytes()
    # Use io.BytesIO to convert bytes back to a file-like object for joblib to load
    vectorizer_buffer = BytesIO(vectorizer_bytes)
    vectorizer = load(vectorizer_buffer)
    
    # Load X_artist
    X_artist = f['X_artist'][:]
    # Load artist names and decode to strings
    artist_names = [name.decode() for name in f['artist_names'][:]]


with h5py.File('conditional_tag_probabilities_matrix.h5', 'r') as f:
    # Reconstruct the sparse co-occurrence matrix
    conditional_co_occurrence_matrix = csr_matrix(
        (f['co_occurrence_data'][:], f['co_occurrence_indices'][:], f['co_occurrence_indptr'][:]),
        shape=f['co_occurrence_shape'][:]
    )

    # Reconstruct the vocabulary
    conditional_words = f['vocabulary_words'][:]
    conditional_indices = f['vocabulary_indices'][:]
    conditional_vocabulary = {key.decode('utf-8'): value for key, value in zip(conditional_words, conditional_indices)}

    # Load the document count
    conditional_doc_count = f['doc_count'][()]
    conditional_smoothing = 100. / conditional_doc_count
    
    
nsfw_tags = set()  # Initialize an empty set to store words meeting the threshold
# Open and read the CSV file
with open("word_rating_probabilities.csv", 'r', newline='', encoding='utf-8') as csvfile:
    reader = csv.reader(csvfile)
    next(reader, None)  # Skip the header row
    for row in reader:
        word = row[0]  # The word is in the first column
        probability_sum = float(row[1])  # The sum of probabilities is in the second column, convert to float for comparison
        # Check if the probability sum meets the threshold and add the word to the set if it does
        if probability_sum >= nsfw_threshold:
            nsfw_tags.add(word)
    
    
# Read the set of valid artists into memory.
artist_set = set()
with open("fluffyrock_3m.csv", 'r', newline='', encoding='utf-8') as csvfile:
    """
    Load artist names from a CSV file and store them in the global set.
    Artist tags start with 'by_' and the prefix will be removed.
    """
    reader = csv.reader(csvfile)
    for row in reader:
        tag_name = row[0]  # Assuming the first column contains the tag names
        if tag_name.startswith('by_'):
            # Strip 'by_' from the start of the tag name and add to the set
            artist_name = tag_name[3:]  # Remove the first three characters 'by_'
            artist_set.add(artist_name)
def is_artist(name):
    return name in artist_set
    
    
sample_images_directory_path = 'sampleimages'
def generate_artist_image_tuples(top_artists, image_directory):
    json_files = glob.glob(f'{image_directory}/*.json')
    json_file_path = json_files[0] if json_files else None
    with open(json_file_path, 'r') as json_file:
        artist_to_file_map = json.load(json_file)
    
    filename = artist_to_file_map.get("")
    image_path = os.path.join(image_directory, filename)
    if os.path.exists(image_path):
        baseline_tuple = [(image_path, "No Artist")]
    
    artist_image_tuples = []
    for artist in top_artists:
        filename = artist_to_file_map.get(artist)
        if filename:
            image_path = os.path.join(image_directory, filename)
            if os.path.exists(image_path):
                artist_image_tuples.append((image_path, artist if artist else "No Artist"))

    return baseline_tuple, artist_image_tuples
    
    
def clean_tag(tag):
    return ''.join(char for char in tag if ord(char) < 128)
    
    
#Normally returns tag to aliases, but when reverse=True, returns alias to tags
def build_aliases_dict(filename, reverse=False):   
    aliases_dict = {}
    with open(filename, 'r', newline='', encoding='utf-8') as csvfile:
        reader = csv.reader(csvfile)
        for row in reader:
            tag = clean_tag(row[0])
            alias_list = [] if row[3] == "null" else [clean_tag(alias) for alias in row[3].split(',')]
            if reverse:
                for alias in alias_list:
                    aliases_dict.setdefault(alias, []).append(tag)
            else:
                aliases_dict[tag] = alias_list
    return aliases_dict
    

def build_tag_count_dict(filename):
    with open(filename, 'r', newline='', encoding='utf-8') as csvfile:
        reader = csv.reader(csvfile)
        result_dict = {}
        for row in reader:
            key = row[0]
            value = int(row[2]) if row[2].isdigit() else None
            if value is not None:
                result_dict[key] = value
    return result_dict

import csv


def build_tag_id_wiki_dict(filename='wiki_pages-2023-08-08.csv'):
    """
    Reads a CSV file and returns a dictionary mapping tag names to tuples of
    (number, most relevant line from the wiki entry). Rows with a non-integer in the first column are ignored.
    The most relevant line is the first line that does not start with "thumb" and is not blank.

    Parameters:
    - filename: The path to the CSV file.

    Returns:
    - A dictionary where each key is a tag name and each value is a tuple (number, most relevant wiki entry line).
    """
    tag_data = {}
    with open(filename, 'r', encoding='utf-8') as csvfile:
        reader = csv.reader(csvfile)

        # Skip the header row
        next(reader)

        for row in reader:
            try:
                # Attempt to convert the first column to an integer
                number = int(row[0])
            except ValueError:
                # If conversion fails, skip this row
                continue

            tag = row[3]
            wiki_entry_full = row[4]

            # Process the wiki_entry to find the most relevant line
            relevant_line = ''
            for line in wiki_entry_full.split('\n'):
                if line.strip() and not line.startswith("thumb"):
                    relevant_line = line
                    break

            # Map the tag to a tuple of (number, relevant_line)
            tag_data[tag] = (number, relevant_line)

    return tag_data


#Imagine we are adding smoothing_value to the number of times word_j occurs in each document for smoothing.
#Note the intention is that sum_i(P(word_i|word_j)) =(approx) # of words in a document rather than 1.
def conditional_probability(word_i, word_j, co_occurrence_matrix, vocabulary, doc_count, smoothing_value=0.01):
    word_i_index = vocabulary.get(word_i)
    word_j_index = vocabulary.get(word_j)
    
    if word_i_index is not None and word_j_index is not None:
        # Directly access the sparse matrix elements
        word_j_count = co_occurrence_matrix[word_j_index, word_j_index]
        smoothed_word_j_count =  word_j_count + (smoothing_value * doc_count)
        
        word_i_count = co_occurrence_matrix[word_i_index, word_i_index]
        
        co_occurrence_count = co_occurrence_matrix[word_i_index, word_j_index]
        smoothed_co_occurrence_count = co_occurrence_count + (smoothing_value * word_i_count) 
        
        # Calculate the conditional probability with smoothing
        conditional_prob = smoothed_co_occurrence_count / smoothed_word_j_count
        
        return conditional_prob
    elif word_i_index is None:
        return 0
    else:
        return None


def geometric_mean_given_words(target_word, context_words, co_occurrence_matrix, vocabulary, doc_count, smoothing_value=0.01):
    probabilities = []
    
    # Collect the conditional probabilities of the target word given each context word, ignoring None values
    for context_word in context_words:
        prob = conditional_probability(target_word, context_word, co_occurrence_matrix, vocabulary, doc_count, smoothing_value)
        if prob is not None:
            probabilities.append(prob)
    
    # Compute the geometric mean of the probabilities, avoiding division by zero
    if probabilities:  # Check if the list is not empty
        geometric_mean = np.prod(probabilities) ** (1.0 / len(probabilities))
    else:
        geometric_mean = 0.5  # Or assign some default value if all probabilities are None
    
    return geometric_mean

    
def create_html_tables_for_tags(subtable_heading, word_similarity_tuples, tag2count, tag2idwiki):
    # Wrap the tag part in a <span> with styles for bold and larger font
    html_str = f"<div style='display: inline-block; margin: 20px; vertical-align: top;'><table><thead><tr><th colspan='3' style='text-align: center; padding-bottom: 10px;'><span style='font-weight: bold; font-size: 20px;'>{subtable_heading}</span></th></tr></thead><tbody><tr style='border-bottom: 1px solid #000;'><th>Corrected Tag</th><th>Similarity</th><th>Count</th></tr>"
    # Loop through the results and add table rows for each
    for word, sim in word_similarity_tuples:
        word_with_underscores = word.replace(' ', '_')
        word_with_escaped_parentheses = word.replace("\\(", "(").replace("\\)", ")").replace("(", "\\(").replace(")", "\\)")
        count = tag2count.get(word_with_underscores.replace("\\(", "(").replace("\\)", ")"), 0)  # Get the count if available, otherwise default to 0
        tag_id, wiki_entry = tag2idwiki.get(word_with_underscores, (None, ''))
        # Check if tag_id and wiki_entry are valid
        if tag_id is not None and wiki_entry:
            # Construct the URL for the tag's wiki page
            wiki_url = f"https://e621.net/wiki_pages/{tag_id}"
            # Make the tag a hyperlink with a tooltip
            tag_element = f"<a href='{wiki_url}' target='_blank' title='{wiki_entry}'>{word_with_escaped_parentheses}</a>"
        else:
            # Display the word without any hyperlink or tooltip
            tag_element = word_with_escaped_parentheses
        # Include the tag element in the table row
        html_str += f"<tr><td style='border: none; padding: 5px; height: 20px;'>{tag_element}</td><td style='border: none; padding: 5px; height: 20px;'>{round(sim, 3)}</td><td style='border: none; padding: 5px; height: 20px;'>{count}</td></tr>"

    html_str += "</tbody></table></div>"
    return html_str
    
    
def create_top_artists_table(top_artists):
    # Add a heading above the table
    html_str = "<div class=\"scrollable-content\" style='display: inline-block; margin: 20px; text-align: center;'>"
    html_str += "<h1>Top Artists</h1>"  # Heading for the table
    # Start the table with increased font size and no borders between rows
    html_str += "<table style='font-size: 20px; border-collapse: collapse;'>"
    html_str += "<thead><tr><th>Artist</th><th>Similarity</th></tr></thead><tbody>"
    # Loop through the top artists and add a row for each without the rank and without borders between rows
    for artist, score in top_artists:
        artist_name = artist[3:] if artist.startswith("by ") else artist  # Remove "by " prefix
        similarity_percentage = "{:.1f}%".format(score * 100)  # Convert score to percentage string with one decimal
        html_str += f"<td style='padding: 3px 20px; border: none;'>{artist_name}</td><td style='padding: 3px 20px; border: none;'>{similarity_percentage}</td></tr>"

    # Close the table HTML
    html_str += "</tbody></table></div>"

    return html_str


def construct_pseudo_vector(pseudo_doc_terms, idf_loaded, tag_to_row_loaded):
    # Initialize a vector of zeros with the length of the term_to_index mapping
    pseudo_vector = np.zeros(len(tag_to_row_loaded))
    
    # Fill in the vector for terms in the pseudo document
    for term in pseudo_doc_terms:
        if term in tag_to_row_loaded:
            index = tag_to_row_loaded[term]
            pseudo_vector[index] = idf_loaded.get(term, 0)
    
    # Return the vector as a 2D array for compatibility with SVD transform
    return pseudo_vector.reshape(1, -1)
    

def get_top_indices(reduced_pseudo_vector, reduced_matrix):
    # Compute cosine similarities
    similarities = cosine_similarity(reduced_pseudo_vector, reduced_matrix).flatten()
    
    # Get sorted tag indices based on similarities, in descending order
    sorted_indices = np.argsort(-similarities)
    
    # Return the top N indices
    return sorted_indices


def get_tfidf_reduced_similar_tags(pseudo_doc_terms, allow_nsfw_tags):
    idf = tf_idf_components['idf']
    term_to_column_index = tf_idf_components['tag_to_column_index']
    row_to_tag = tf_idf_components['row_to_tag']
    reduced_matrix = tf_idf_components['reduced_matrix']
    svd = tf_idf_components['svd_model']

    # Construct the TF-IDF vector
    pseudo_tfidf_vector = construct_pseudo_vector(pseudo_doc_terms, idf, term_to_column_index)

    # Reduce the dimensionality of the pseudo-document vector for the reduced matrix
    reduced_pseudo_vector = svd.transform(pseudo_tfidf_vector)
    
    # Compute cosine similarities in the reduced space
    cosine_similarities_reduced = cosine_similarity(reduced_pseudo_vector, reduced_matrix).flatten()

    # Sort the indices by descending cosine similarity
    top_indices_reduced = np.argsort(cosine_similarities_reduced)
    
    # Map indices to tags with their similarities
    tag_similarity_dict = {row_to_tag[i]: cosine_similarities_reduced[i] for i in top_indices_reduced if i in row_to_tag}

    if not allow_nsfw_tags:
        tag_similarity_dict = {tag: sim for tag, sim in tag_similarity_dict.items() if tag not in nsfw_tags}
        
    tag_similarity_dict = {"by " + tag if is_artist(tag) else tag: sim for tag, sim in tag_similarity_dict.items()}

    # Sort and transform tag names
    sorted_tag_similarity_dict = OrderedDict(sorted(tag_similarity_dict.items(), key=lambda x: x[1], reverse=True))
    transformed_sorted_tag_similarity_dict = OrderedDict(
        (key.replace('_', ' ').replace('(', '\\(').replace(')', '\\)'), value)
        for key, value in sorted_tag_similarity_dict.items()
    )

    return transformed_sorted_tag_similarity_dict


def create_html_placeholder(title="", content="", placeholder_height=400, placeholder_width="100%"):
    # Include a title in the same style as the top artists table heading
    html_placeholder = f"<div class=\"scrollable-content\" style='text-align: center;'><h1>{title}</h1></div>"
    # Conditionally add content if present
    if content:
        html_placeholder += f"<div style='text-align: center; margin-bottom: 20px;'><p>{content}</p></div>"
    # Add the placeholder div with specified height and width
    html_placeholder += f"<div style='height: {placeholder_height}px; width: {placeholder_width}; margin: 20px auto; background: transparent;'></div>"
    return html_placeholder
    

def find_similar_tags(test_tags, similarity_weight, allow_nsfw_tags):
    #Initialize stuff
    if not hasattr(find_similar_tags, "fasttext_small_model"):
        find_similar_tags.fasttext_small_model = compress_fasttext.models.CompressedFastTextKeyedVectors.load('e621FastTextModel010Replacement_small.bin')
    tag_aliases_file = 'fluffyrock_3m.csv'
    if not hasattr(find_similar_tags, "tag2aliases"):
        find_similar_tags.tag2aliases = build_aliases_dict(tag_aliases_file)
    if not hasattr(find_similar_tags, "alias2tags"):
        find_similar_tags.alias2tags = build_aliases_dict(tag_aliases_file, reverse=True)
    if not hasattr(find_similar_tags, "tag2count"):
        find_similar_tags.tag2count = build_tag_count_dict(tag_aliases_file)
    if not hasattr(find_similar_tags, "tag2idwiki"):    
        find_similar_tags.tag2idwiki = build_tag_id_wiki_dict()
    
    modified_tags = [tag_info['modified_tag'] for tag_info in test_tags]
    transformed_tags = [tag.replace(' ', '_') for tag in modified_tags]

    # Find similar tags and prepare data for tables
    html_content = "<div class=\"scrollable-content\" style='display: inline-block; margin: 20px; text-align: center;'>"
    html_content += "<h1>Unknown Tags</h1>"  # Heading for the table
    tags_added = False
    bad_entities = []
    encountered_modified_tags = set()
    for tag_info in test_tags:
        original_tag = tag_info['original_tag']
        modified_tag = tag_info['modified_tag']
        start_pos = tag_info['start_pos']
        end_pos = tag_info['end_pos']
        node_type = tag_info['node_type']
        
        if modified_tag in special_tags:
            bad_entities.append({"entity":"Special", "start":start_pos, "end":end_pos})
            continue
            
        if modified_tag in encountered_modified_tags:
            bad_entities.append({"entity":"Duplicate", "start":start_pos, "end":end_pos})
            continue
        encountered_modified_tags.add(modified_tag)
        
        if node_type == "double_comma":
            bad_entities.append({"entity":"Double Comma", "start":start_pos, "end":end_pos})
            continue
            
        modified_tag_for_search = modified_tag.replace(' ','_')
        similar_words = find_similar_tags.fasttext_small_model.most_similar(modified_tag_for_search, topn = 100)
        result, seen = [], set(transformed_tags)
        
        if modified_tag_for_search in find_similar_tags.tag2aliases:
            if modified_tag in find_similar_tags.tag2aliases and "_" in modified_tag:   #Implicitly tell the user that they should get rid of the underscore
                result.append(modified_tag_for_search.replace('_',' '), 1)
                seen.add(modified_tag)
            else:   #The user correctly did not put underscores in their tag
                continue
        else:
            for item in similar_words:
                similar_word, similarity = item
                if similar_word not in seen:
                    if similar_word in find_similar_tags.tag2aliases:
                        result.append((similar_word.replace('_', ' '), round(similarity, 3)))
                        seen.add(similar_word)
                    else:
                        for similar_tag in find_similar_tags.alias2tags.get(similar_word, []):
                            if similar_tag not in seen:
                                result.append((similar_tag.replace('_', ' '), round(similarity, 3)))
                                seen.add(similar_tag)

        #Remove NSFW tags if appropriate.
        if not allow_nsfw_tags:
            result = [(word, score) for word, score in result if word.replace(' ','_') not in nsfw_tags]

        #Adjust score based on context
        for i in range(len(result)):
            word, score = result[i]  # Unpack the tuple
            geometric_mean = geometric_mean_given_words(word.replace(' ','_'), [context_tag for context_tag in transformed_tags if context_tag != word and context_tag != modified_tag], conditional_co_occurrence_matrix, conditional_vocabulary, conditional_doc_count, smoothing_value=conditional_smoothing)
            adjusted_score = (similarity_weight * geometric_mean) + ((1-similarity_weight)*score)  # Apply the adjustment function
            result[i] = (word, adjusted_score)  # Update the tuple with the adjusted score
            #print(word, score, geometric_mean, adjusted_score)

        result = sorted(result, key=lambda x: x[1], reverse=True)[:10]
        html_content += create_html_tables_for_tags(modified_tag, result, find_similar_tags.tag2count, find_similar_tags.tag2idwiki)
        
        bad_entities.append({"entity":"Unknown Tag", "start":start_pos, "end":end_pos})
        
        tags_added=True
    # If no tags were processed, add a message
    if not tags_added:
        html_content = create_html_placeholder(title="Unknown Tags", content="No Unknown Tags Found")

    return html_content, bad_entities  # Return list of lists for Dataframe


def build_tag_offsets_dicts(new_image_tags_with_positions):
    # Structure the data for HighlightedText
    tag_data = []
    for tag_text, start_pos, nodetype in new_image_tags_with_positions:
        # Modify the tag
        modified_tag = tag_text.replace('_', ' ').replace('\\(', '(').replace('\\)', ')').strip()
        artist_matrix_tag = tag_text.replace('_', ' ').replace('\\(', '\(').replace('\\)', '\)').strip()
        tf_idf_matrix_tag = re.sub(r'\\([()])', r'\1', re.sub(r' ', '_', tag_text.strip().removeprefix('by ').removeprefix('by_')))
        # Calculate the end position based on the original tag length
        end_pos = start_pos + len(tag_text)
        # Append the structured data for each tag
        tag_data.append({
            "original_tag": tag_text,
            "start_pos": start_pos,
            "end_pos": end_pos,
            "modified_tag": modified_tag,
            "artist_matrix_tag": artist_matrix_tag,
            "tf_idf_matrix_tag": tf_idf_matrix_tag, 
            "node_type": nodetype
        })
    return tag_data
    
    
def augment_bad_entities_with_regex(text):
    bad_entities = []

    #comma at end
    match = re.search(r',(?=\s*$)', text)
    if match:
        index = match.start()
        bad_entities.append({"entity":"Remove Final Comma", "start":index, "end":index+1})
    match = re.search(r'\([^()]*(,)\s*\)\s*$', text)
    if match:
        index = match.start(1)
        bad_entities.append({"entity":"Remove Final Comma", "start":index, "end":index+1})
    match = re.search(r'\([^()]*(,)\s*:\s*\d+(\.\d+)?\s*\)\s*$', text)
    if match:
        index = match.start(1)
        bad_entities.append({"entity":"Remove Final Comma", "start":index, "end":index+1})

   # Comma after parentheses, multiple occurrences
    for match in re.finditer(r'\)\s*(,)\s*[^\s]', text):
        index = match.start(1)
        bad_entities.append({"entity": "Move Comma Inside Parentheses", "start": index, "end": index + 1})

    return bad_entities
    

def find_similar_artists(original_tags_string, top_n, similarity_weight, allow_nsfw_tags):
    try:
        new_tags_string = original_tags_string.lower()
        new_tags_string, removed_tags = remove_special_tags(new_tags_string)
        
        # Parse the prompt
        parsed = parser.parse(new_tags_string)
        
        # Extract tags from the parsed tree
        new_image_tags = extract_tags(parsed)
        
        tag_data = build_tag_offsets_dicts(new_image_tags)

        ###unseen_tags = list(set(OrderedDict.fromkeys(new_image_tags)) - set(vectorizer.vocabulary_.keys()))   #We may want this line again later.  These are the tags that were not used to calculate the artists list.
        unseen_tags_data, bad_entities = find_similar_tags(tag_data, similarity_weight, allow_nsfw_tags)
    
        #Bad tags stuff
        bad_entities.extend(augment_bad_entities_with_regex(new_tags_string))
        bad_entities.sort(key=lambda x: x['start'])
        bad_tags_illustrated_string = {"text":new_tags_string, "entities":bad_entities}

        #Suggested tags stuff
        suggested_tags_html_content = "<div class=\"scrollable-content\" style='display: inline-block; margin: 20px; text-align: center;'>"

        suggested_tags_html_content += "<h1>Suggested Tags</h1>"  # Heading for the table
        suggested_tags = get_tfidf_reduced_similar_tags([item["tf_idf_matrix_tag"] for item in tag_data] + removed_tags, allow_nsfw_tags)
        
        # Create a set of tags that should be filtered out
        filter_tags = {entry["original_tag"].strip() for entry in tag_data}
        # Use this set to filter suggested_tags
        suggested_tags_filtered = OrderedDict((k, v) for k, v in suggested_tags.items() if k not in filter_tags)
        
        # Splitting the dictionary into two based on the condition
        suggested_artist_tags_filtered = OrderedDict((k, v) for k, v in suggested_tags_filtered.items() if k.startswith("by "))
        suggested_non_artist_tags_filtered = OrderedDict((k, v) for k, v in suggested_tags_filtered.items() if not k.startswith("by ") and k not in special_tags)

        topnsuggestions = list(islice(suggested_non_artist_tags_filtered.items(), 100))
        suggested_tags_html_content += create_html_tables_for_tags("Suggested Tag", topnsuggestions, find_similar_tags.tag2count, find_similar_tags.tag2idwiki)

        #Artist stuff
        #artist_matrix_tags = [tag_info['artist_matrix_tag'] for tag_info in tag_data if tag_info['node_type'] == "tag"]
        #X_new_image = vectorizer.transform([','.join(artist_matrix_tags + removed_tags)])
        #similarities = cosine_similarity(X_new_image, X_artist)[0]
        #
        #top_artist_indices = np.argsort(similarities)[-(top_n + 1):][::-1]
        #top_artists = [(artist_names[i], similarities[i]) for i in top_artist_indices if artist_names[i].lower() != "by conditional dnp"][:top_n]

        excluded_artists = ["by conditional dnp", "by unknown artist"]
        top_artists = [(key, value) for key, value in  suggested_artist_tags_filtered.items() if key.lower() not in excluded_artists][:top_n]
        top_artists_str = create_top_artists_table(top_artists)
        dynamic_prompts_formatted_artists = "{" + "|".join([artist for artist, _ in top_artists]) + "}"

        image_galleries = []
        for root, dirs, files in os.walk(sample_images_directory_path):
            for name in dirs:    
                baseline, artists = generate_artist_image_tuples([name[3:] for name, _ in top_artists], os.path.join(root, name))
                image_galleries.append(baseline)  # Add baseline as its own gallery item
                image_galleries.append(artists)  # Extend the list with artist tuples

        return (unseen_tags_data, bad_tags_illustrated_string, suggested_tags_html_content, top_artists_str, dynamic_prompts_formatted_artists, *image_galleries)
    except ParseError as e:
        return [], "Parse Error: Check for mismatched parentheses or something", "", "", None, None


with gr.Blocks(css=css) as app:
    with gr.Group():
        with gr.Row():
            with gr.Column(scale=3):
                image_tags = gr.Textbox(label="Enter Prompt", placeholder="e.g. fox, outside, detailed background, ...")
                bad_tags_illustrated_string = gr.HighlightedText(show_legend=True, color_map={"Unknown Tag":"red","Duplicate":"yellow","Remove Final Comma":"purple","Move Comma Inside Parentheses":"green"}, label="Annotated Prompt")
            with gr.Column(scale=1):
                #image_path = os.path.join("https://huggingface.co/spaces/FoodDesert/Prompt_Squirrel/resolve/main", "transparentsquirrel.png")
                #gr.Image(label=" ", value=image_path, height=155, width=140)
                gr.HTML('<div style="text-align: center;"><img src="https://huggingface.co/spaces/FoodDesert/Prompt_Squirrel/resolve/main/mascotimages/transparentsquirrel.png" alt="Cute Mascot" style="max-height: 180px; background: transparent;"></div><br>')
                #gr.HTML("<br>" * 2)  # Adjust the number of line breaks ("<br>") as needed to push the button down
                #image_path = os.path.join('mascotimages', "transparentsquirrel.png") 
                #random_image_path = os.path.join('mascotimages', random.choice([f for f in os.listdir('mascotimages') if os.path.isfile(os.path.join('mascotimages', f))]))
                #with Image.open(random_image_path) as img:
                #    gr.Image(value=img,show_label=False, show_download_button=False, show_share_button=False, height=200)
                #gr.Image(value="https://huggingface.co/spaces/FoodDesert/Prompt_Squirrel/resolve/main/mascotimages/transparentsquirrel.png",show_label=False, show_download_button=False, show_share_button=False, height=200)
                #I posted the image to discord, and that's where this link came from.  This is a very ugly way to do this, but I could not, no matter what I tried, get it to display an image from within the space itself.  The galleries work fine for some reason, but not this.
                #gr.Image(value="https://res.cloudinary.com/dnse84ol6/image/upload/v1713538125/transparentsquirrel_zhou7f.png",show_label=False, show_download_button=False, show_share_button=False, height=200)
                submit_button = gr.Button(variant="primary")
    with gr.Row():
        with gr.Column(scale=3):    
            with gr.Group():
                with gr.Row():
                    similarity_weight = gr.Slider(minimum=0, maximum=1, value=0.5, step=0.1, label="Similarity weight")
                    allow_nsfw = gr.Checkbox(label="Allow NSFW Tags", value=False)
                with gr.Row():
                    with gr.Column(scale=2):
                        unseen_tags = gr.HTML(label="Unknown Tags", value=create_html_placeholder(title="Unknown Tags"))
                    with gr.Column(scale=1):
                        suggested_tags = gr.HTML(label="Suggested Tags", value=create_html_placeholder(title="Suggested Tags"))
        with gr.Column(scale=1):
            with gr.Group():
                num_artists = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Number of artists")
                top_artists = gr.HTML(label="Top Artists", value=create_html_placeholder(title="Top Artists"))
                dynamic_prompts = gr.Textbox(label="Dynamic Prompts Format", info="For if you're using the Automatic1111 webui (https://github.com/AUTOMATIC1111/stable-diffusion-webui) with the Dynamic Prompts extension activated (https://github.com/adieyal/sd-dynamic-prompts) and want to try them all individually.")
    galleries = []
    for root, dirs, files in os.walk(sample_images_directory_path):
        for name in dirs:
            with gr.Row():
                baseline = gr.Gallery(allow_preview=False, rows=1, columns=1, height=420, scale=3)
                styles = gr.Gallery(preview=False, rows=2, columns=5, height=420, scale=8)
                galleries.extend([baseline, styles])

    submit_button.click(
        find_similar_artists, 
        inputs=[image_tags, num_artists, similarity_weight, allow_nsfw], 
        outputs=[unseen_tags, bad_tags_illustrated_string, suggested_tags, top_artists, dynamic_prompts] + galleries
    )
    
    gr.Markdown(faq_content)
    

app.launch()