[CLEANUP][Pegasus]

pull/286/head
Kye 1 year ago
parent 277c97411f
commit 815b4141be

@ -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
Loading…
Cancel
Save