WibddlZddlmZddlmZddlmZmZddlm Z ddl m Z Gdde Z y) N)Image)tqdm)ListUnion)SentenceTransformer) BaseEmbedderc eZdZdZ ddeeefdeeefdeffd Zdde ede edzd e d e jfd Z dde ed e d e jfd Zdd e ed e d e jfdZdZdZxZS)MultiModalBackendaMultimodal backend using Sentence-transformers. The sentence-transformers embedding model used for generating word, document, and image embeddings. Arguments: embedding_model: A sentence-transformers embedding model that can either embed both images and text or only text. If it only embeds text, then `image_model` needs to be used to embed the images. image_model: A sentence-transformers embedding model that is used to embed only images. batch_size: The sizes of image batches to pass Examples: To create a model, you can load in a string pointing to a sentence-transformers model: ```python from bertopic.backend import MultiModalBackend sentence_model = MultiModalBackend("clip-ViT-B-32") ``` or you can instantiate a model yourself: ```python from bertopic.backend import MultiModalBackend from sentence_transformers import SentenceTransformer embedding_model = SentenceTransformer("clip-ViT-B-32") sentence_model = MultiModalBackend(embedding_model) ``` Nembedding_model image_model batch_sizec t|||_t|tr||_n,t|t rt ||_n tdd|_|Dt|tr||_n,t|t rt ||_n td |j jjj|_ y#t$r|j j|_ Yyd|_ YyxYw)NzPlease select a correct SentenceTransformers model: `from sentence_transformers import SentenceTransformer` `model = SentenceTransformer('clip-ViT-B-32')`) super__init__r isinstancerr str ValueErrorr _first_module processor tokenizerAttributeError)selfr r r __class__s h/home/sietch6/trending-topics-pipeline/venv/lib/python3.12/site-packages/bertopic/backend/_multimodal.pyrzMultiModalBackend.__init__-s $ o': ;#2D  -#6#GD A    "+':;#. K-#6{#C  E  "!11??AKKUUDN <!11;;DN "!DNs)3C$D D  documentsimagesverbosereturncd}|d|j|}d}t|tr|j||}d}||t j ||gd}||S||S||Sy)aGEmbed a list of n documents/words or images into an n-dimensional matrix of embeddings. Either documents, images, or both can be provided. If both are provided, then the embeddings are averaged. Arguments: documents: A list of documents or words to be embedded images: A list of image paths to be embedded verbose: Controls the verbosity of the process Returns: Document/words embeddings with shape (n, m) with `n` documents/words that each have an embeddings size of `m` Nr)axis)embed_documentsrlist embed_imagesnpmean)rrrrdoc_embeddingsimage_embeddingsaveraged_embeddingss rembedzMultiModalBackend.embedWs" Q< #!11);K*LST"U   *& &  '! !  )# #*c|Dcgc]}|j|}}|jj||}|Scc}w)aEmbed a list of n documents/words into an n-dimensional matrix of embeddings. Arguments: documents: A list of documents or words to be embedded verbose: Controls the verbosity of the process Returns: Document/words embeddings with shape (n, m) with `n` documents/words that each have an embeddings size of `m` show_progress_bar)_truncate_documentr encode)rrrdoctruncated_docs embeddingss rr!z!MultiModalBackend.embed_documents}sLCLL3$11#6LL))00SZ0[ Ms>wordsc@|jj||}|S)amEmbed a list of n words into an n-dimensional matrix of embeddings. Arguments: words: A list of words to be embedded verbose: Controls the verbosity of the process Returns: Document/words embeddings with shape (n, m) with `n` documents/words that each have an embeddings size of `m` r,)r r/)rr3rr2s r embed_wordszMultiModalBackend.embed_wordss%))00'0R r*c|jrZttjt ||jz }g}t t || D]}||jz}||jz|jz}|||Dcgc])}t|trtj|n|+} }|j|jj| } n|jj| d} |j| jt|dts| D]}|j!tj"|}|S|D cgc]} tj| } } |j|jj| }|S|jj| d}|Scc}wcc} w)N)disableFr,r)r intr$ceillenrrangerrropenr r/r extendtolistclosearray) rrr nr_iterationsr2i start_index end_indeximageimages_to_embedimg_embfilepaths rr#zMultiModalBackend.embed_imagess ??F doo(E FGMJ% .G D &$//1 0DOOC Y__jktXu#OTE3)?EJJu%UJ####/"..55oFG"2299/]b9cG!!'.."23fQi-!0& & &"*-J EKKuzz(3KOK+!--44_E "1188\a8b )#Ls .G,Gc|jrZ|jj|}t|dkDr1|dd}|jj|}|j |S|S)NML)rr/r:decoder.)rdocumenttokenstruncated_tokenss rr.z$MultiModalBackend._truncate_documentsb >>^^**84F6{R#)!B< >>001AB..x88r*)N )NF)F)__name__ __module__ __qualname____doc__rrrr8rrboolr$ndarrayr)r!r5r#r. __classcell__)rs@rr r s J8< ("s$778("3 334(" ("T$$tCy$$$s)d2B$$TX$$egeoeo$$LcTbjj c T bjj @ r*r ) numpyr$PILrrtypingrrsentence_transformersrbertopic.backendrr r*rr_s%5)~ ~r*