From 815b4141be8b406f3a42b9d81c282c360c863150 Mon Sep 17 00:00:00 2001 From: Kye Date: Sat, 9 Dec 2023 22:59:20 -0800 Subject: [PATCH] [CLEANUP][Pegasus] --- swarms/models/pegasus.py | 64 ---------------------------------------- 1 file changed, 64 deletions(-) delete mode 100644 swarms/models/pegasus.py diff --git a/swarms/models/pegasus.py b/swarms/models/pegasus.py deleted file mode 100644 index b2f43980..00000000 --- a/swarms/models/pegasus.py +++ /dev/null @@ -1,64 +0,0 @@ -import logging -from typing import Union - -from pegasus import Pegasus - - -class PegasusEmbedding: - """ - Pegasus - - Args: - modality (str): Modality to use for embedding - multi_process (bool, optional): Whether to use multi-process. Defaults to False. - n_processes (int, optional): Number of processes to use. Defaults to 4. - - Usage: - -------------- - pegasus = PegasusEmbedding(modality="text") - pegasus.embed("Hello world") - - - vision - -------------- - pegasus = PegasusEmbedding(modality="vision") - pegasus.embed("https://i.imgur.com/1qZ0K8r.jpeg") - - audio - -------------- - pegasus = PegasusEmbedding(modality="audio") - pegasus.embed("https://www2.cs.uic.edu/~i101/SoundFiles/StarWars60.wav") - - - - """ - - def __init__( - self, - modality: str, - multi_process: bool = False, - n_processes: int = 4, - ): - self.modality = modality - self.multi_process = multi_process - self.n_processes = n_processes - try: - self.pegasus = Pegasus( - modality, multi_process, n_processes - ) - except Exception as e: - logging.error( - "Failed to initialize Pegasus with modality:" - f" {modality}: {e}" - ) - raise - - def embed(self, data: Union[str, list[str]]): - """Embed the data""" - try: - return self.pegasus.embed(data) - except Exception as e: - logging.error( - f"Failed to generate embeddings. Error: {e}" - ) - raise