Kye
11 months ago
parent
aff03cb649
commit
1550d8c7ff
@ -1,227 +0,0 @@
|
||||
r"""
|
||||
BioGPT
|
||||
Pre-trained language models have attracted increasing attention in the biomedical domain,
|
||||
inspired by their great success in the general natural language domain.
|
||||
Among the two main branches of pre-trained language models in the general language domain, i.e. BERT (and its variants) and GPT (and its variants),
|
||||
the first one has been extensively studied in the biomedical domain, such as BioBERT and PubMedBERT.
|
||||
While they have achieved great success on a variety of discriminative downstream biomedical tasks,
|
||||
the lack of generation ability constrains their application scope.
|
||||
In this paper, we propose BioGPT, a domain-specific generative Transformer language model
|
||||
pre-trained on large-scale biomedical literature.
|
||||
We evaluate BioGPT on six biomedical natural language processing tasks
|
||||
and demonstrate that our model outperforms previous models on most tasks.
|
||||
Especially, we get 44.98%, 38.42% and 40.76% F1 score on BC5CDR, KD-DTI and DDI
|
||||
end-to-end relation extraction tasks, respectively, and 78.2% accuracy on PubMedQA,
|
||||
creating a new record. Our case study on text generation further demonstrates the
|
||||
advantage of BioGPT on biomedical literature to generate fluent descriptions for biomedical terms.
|
||||
|
||||
|
||||
@article{10.1093/bib/bbac409,
|
||||
author = {Luo, Renqian and Sun, Liai and Xia, Yingce and Qin, Tao and Zhang, Sheng and Poon, Hoifung and Liu, Tie-Yan},
|
||||
title = "{BioGPT: generative pre-trained transformer for biomedical text generation and mining}",
|
||||
journal = {Briefings in Bioinformatics},
|
||||
volume = {23},
|
||||
number = {6},
|
||||
year = {2022},
|
||||
month = {09},
|
||||
abstract = "{Pre-trained language models have attracted increasing attention in the biomedical domain, inspired by their great success in the general natural language domain. Among the two main branches of pre-trained language models in the general language domain, i.e. BERT (and its variants) and GPT (and its variants), the first one has been extensively studied in the biomedical domain, such as BioBERT and PubMedBERT. While they have achieved great success on a variety of discriminative downstream biomedical tasks, the lack of generation ability constrains their application scope. In this paper, we propose BioGPT, a domain-specific generative Transformer language model pre-trained on large-scale biomedical literature. We evaluate BioGPT on six biomedical natural language processing tasks and demonstrate that our model outperforms previous models on most tasks. Especially, we get 44.98\%, 38.42\% and 40.76\% F1 score on BC5CDR, KD-DTI and DDI end-to-end relation extraction tasks, respectively, and 78.2\% accuracy on PubMedQA, creating a new record. Our case study on text generation further demonstrates the advantage of BioGPT on biomedical literature to generate fluent descriptions for biomedical terms.}",
|
||||
issn = {1477-4054},
|
||||
doi = {10.1093/bib/bbac409},
|
||||
url = {https://doi.org/10.1093/bib/bbac409},
|
||||
note = {bbac409},
|
||||
eprint = {https://academic.oup.com/bib/article-pdf/23/6/bbac409/47144271/bbac409.pdf},
|
||||
}
|
||||
"""
|
||||
|
||||
import torch
|
||||
from transformers import (
|
||||
BioGptForCausalLM,
|
||||
BioGptTokenizer,
|
||||
pipeline,
|
||||
set_seed,
|
||||
)
|
||||
|
||||
|
||||
class BioGPT:
|
||||
"""
|
||||
A wrapper class for the BioGptForCausalLM model from the transformers library.
|
||||
|
||||
Attributes:
|
||||
model_name (str): Name of the pretrained model.
|
||||
model (BioGptForCausalLM): The pretrained BioGptForCausalLM model.
|
||||
tokenizer (BioGptTokenizer): The tokenizer for the BioGptForCausalLM model.
|
||||
|
||||
Methods:
|
||||
__call__: Generate text based on the given input.
|
||||
get_features: Get the features of a given text.
|
||||
beam_search_decoding: Generate text using beam search decoding.
|
||||
set_pretrained_model: Set a new tokenizer and model.
|
||||
get_config: Get the model's configuration.
|
||||
save_model: Save the model and tokenizer to a directory.
|
||||
load_from_path: Load a model and tokenizer from a directory.
|
||||
print_model: Print the model's architecture.
|
||||
|
||||
Usage:
|
||||
>>> from swarms.models.biogpt import BioGPTWrapper
|
||||
>>> model = BioGPTWrapper()
|
||||
>>> out = model("The patient has a fever")
|
||||
>>> print(out)
|
||||
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str = "microsoft/biogpt",
|
||||
max_length: int = 500,
|
||||
num_return_sequences: int = 5,
|
||||
do_sample: bool = True,
|
||||
min_length: int = 100,
|
||||
):
|
||||
"""
|
||||
Initialize the wrapper class with a model name.
|
||||
|
||||
Args:
|
||||
model_name (str): Name of the pretrained model. Default is "microsoft/biogpt".
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self.max_length = max_length
|
||||
self.num_return_sequences = num_return_sequences
|
||||
self.do_sample = do_sample
|
||||
self.min_length = min_length
|
||||
|
||||
self.model = BioGptForCausalLM.from_pretrained(
|
||||
self.model_name
|
||||
)
|
||||
self.tokenizer = BioGptTokenizer.from_pretrained(
|
||||
self.model_name
|
||||
)
|
||||
|
||||
def __call__(self, text: str):
|
||||
"""
|
||||
Generate text based on the given input.
|
||||
|
||||
Args:
|
||||
text (str): The input text to generate from.
|
||||
max_length (int): Maximum length of the generated text.
|
||||
num_return_sequences (int): Number of sequences to return.
|
||||
do_sample (bool): Whether or not to use sampling in generation.
|
||||
|
||||
Returns:
|
||||
list[dict]: A list of generated texts.
|
||||
"""
|
||||
set_seed(42)
|
||||
generator = pipeline(
|
||||
"text-generation",
|
||||
model=self.model,
|
||||
tokenizer=self.tokenizer,
|
||||
)
|
||||
out = generator(
|
||||
text,
|
||||
max_length=self.max_length,
|
||||
num_return_sequences=self.num_return_sequences,
|
||||
do_sample=self.do_sample,
|
||||
)
|
||||
|
||||
return out[0]["generated_text"]
|
||||
|
||||
def get_features(self, text):
|
||||
"""
|
||||
Get the features of a given text.
|
||||
|
||||
Args:
|
||||
text (str): Input text.
|
||||
|
||||
Returns:
|
||||
BaseModelOutputWithPastAndCrossAttentions: Model output.
|
||||
"""
|
||||
encoded_input = self.tokenizer(text, return_tensors="pt")
|
||||
return self.model(**encoded_input)
|
||||
|
||||
def beam_search_decoding(
|
||||
self,
|
||||
sentence,
|
||||
num_beams=5,
|
||||
early_stopping=True,
|
||||
):
|
||||
"""
|
||||
Generate text using beam search decoding.
|
||||
|
||||
Args:
|
||||
sentence (str): The input sentence to generate from.
|
||||
min_length (int): Minimum length of the generated text.
|
||||
max_length (int): Maximum length of the generated text.
|
||||
num_beams (int): Number of beams for beam search.
|
||||
early_stopping (bool): Whether to stop early during beam search.
|
||||
|
||||
Returns:
|
||||
str: The generated text.
|
||||
"""
|
||||
inputs = self.tokenizer(sentence, return_tensors="pt")
|
||||
set_seed(42)
|
||||
with torch.no_grad():
|
||||
beam_output = self.model.generate(
|
||||
**inputs,
|
||||
min_length=self.min_length,
|
||||
max_length=self.max_length,
|
||||
num_beams=num_beams,
|
||||
early_stopping=early_stopping,
|
||||
)
|
||||
return self.tokenizer.decode(
|
||||
beam_output[0], skip_special_tokens=True
|
||||
)
|
||||
|
||||
# Feature 1: Set a new tokenizer and model
|
||||
def set_pretrained_model(self, model_name):
|
||||
"""
|
||||
Set a new tokenizer and model.
|
||||
|
||||
Args:
|
||||
model_name (str): Name of the pretrained model.
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self.model = BioGptForCausalLM.from_pretrained(
|
||||
self.model_name
|
||||
)
|
||||
self.tokenizer = BioGptTokenizer.from_pretrained(
|
||||
self.model_name
|
||||
)
|
||||
|
||||
# Feature 2: Get the model's config details
|
||||
def get_config(self):
|
||||
"""
|
||||
Get the model's configuration.
|
||||
|
||||
Returns:
|
||||
PretrainedConfig: The configuration of the model.
|
||||
"""
|
||||
return self.model.config
|
||||
|
||||
# Feature 3: Save the model and tokenizer to disk
|
||||
def save_model(self, path):
|
||||
"""
|
||||
Save the model and tokenizer to a directory.
|
||||
|
||||
Args:
|
||||
path (str): Path to the directory.
|
||||
"""
|
||||
self.model.save_pretrained(path)
|
||||
self.tokenizer.save_pretrained(path)
|
||||
|
||||
# Feature 4: Load a model from a custom path
|
||||
def load_from_path(self, path):
|
||||
"""
|
||||
Load a model and tokenizer from a directory.
|
||||
|
||||
Args:
|
||||
path (str): Path to the directory.
|
||||
"""
|
||||
self.model = BioGptForCausalLM.from_pretrained(path)
|
||||
self.tokenizer = BioGptTokenizer.from_pretrained(path)
|
||||
|
||||
# Feature 5: Print the model's architecture
|
||||
def print_model(self):
|
||||
"""
|
||||
Print the model's architecture.
|
||||
"""
|
||||
print(self.model)
|
||||
@ -1,91 +0,0 @@
|
||||
import io
|
||||
|
||||
import requests
|
||||
import torch
|
||||
from PIL import Image
|
||||
from transformers import (
|
||||
AutoModelForCausalLM,
|
||||
AutoProcessor,
|
||||
GenerationConfig,
|
||||
)
|
||||
|
||||
from swarms.models.base_multimodal_model import (
|
||||
BaseMultiModalModel,
|
||||
) # noqa: F401
|
||||
|
||||
|
||||
class ChestMultiModalAgent(BaseMultiModalModel):
|
||||
"""
|
||||
Initialize the ChestAgent.
|
||||
|
||||
Args:
|
||||
device (str): The device to run the model on. Default is "cuda".
|
||||
dtype (torch.dtype): The data type to use for the model. Default is torch.float16.
|
||||
model_name (str): The name or path of the pre-trained model to use. Default is "StanfordAIMI/CheXagent-8b".
|
||||
|
||||
Example:
|
||||
>>> agent = ChestAgent()
|
||||
>>> agent.run("What are the symptoms of COVID-19?", "https://example.com/image.jpg")
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
device="cuda",
|
||||
dtype=torch.float16,
|
||||
model_name="StanfordAIMI/CheXagent-8b",
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
# Step 1: Setup constants
|
||||
self.device = device
|
||||
self.dtype = dtype
|
||||
|
||||
# Step 2: Load Processor and Model
|
||||
self.processor = AutoProcessor.from_pretrained(
|
||||
model_name, trust_remote_code=True
|
||||
)
|
||||
self.generation_config = GenerationConfig.from_pretrained(
|
||||
model_name
|
||||
)
|
||||
self.model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name,
|
||||
torch_dtype=self.dtype,
|
||||
trust_remote_code=True,
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def run(self, task: str, img: str, *args, **kwargs):
|
||||
"""
|
||||
Run the ChestAgent to generate findings based on an image and a prompt.
|
||||
|
||||
Args:
|
||||
image_path (str): The URL or local path of the image.
|
||||
prompt (str): The prompt to use for generating findings.
|
||||
|
||||
Returns:
|
||||
str: The generated findings.
|
||||
"""
|
||||
# Step 3: Fetch the images
|
||||
images = [
|
||||
Image.open(io.BytesIO(requests.get(img).content)).convert(
|
||||
"RGB"
|
||||
)
|
||||
]
|
||||
|
||||
# Step 4: Generate the Findings section
|
||||
inputs = self.processor(
|
||||
images=images,
|
||||
text=f" USER: <s>{task} ASSISTANT: <s>",
|
||||
return_tensors="pt",
|
||||
).to(device=self.device, dtype=self.dtype)
|
||||
output = self.model.generate(
|
||||
**inputs,
|
||||
generation_config=self.generation_config,
|
||||
)[0]
|
||||
response = self.processor.tokenizer.decode(
|
||||
output, skip_special_tokens=True
|
||||
)
|
||||
|
||||
return response
|
||||
@ -1,180 +0,0 @@
|
||||
from io import BytesIO
|
||||
|
||||
import requests
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torchvision.transforms import GaussianBlur
|
||||
from transformers import CLIPModel, CLIPProcessor
|
||||
|
||||
|
||||
class CLIPQ:
|
||||
"""CLIPQ model for image and text retrieval
|
||||
|
||||
Args:
|
||||
model_name (str): The name of the CLIP model to use
|
||||
query_text (str): The query text to use for the model
|
||||
|
||||
Example:
|
||||
>>> clipq = CLIPQ()
|
||||
>>> image = clipq.fetch_image_from_url()
|
||||
>>> vectors = clipq.get_vectors(image)
|
||||
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str = "openai/clip-vit-base-patch16",
|
||||
query_text: str = "A photo ",
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
self.model = CLIPModel.from_pretrained(
|
||||
model_name, *args, **kwargs
|
||||
)
|
||||
self.processor = CLIPProcessor.from_pretrained(model_name)
|
||||
self.query_text = query_text
|
||||
|
||||
def fetch_image_from_url(self, url="https://picsum.photos/800"):
|
||||
"""Fetches an image from the given url"""
|
||||
response = requests.get(url)
|
||||
if response.status_code != 200:
|
||||
raise Exception("Failed to fetch an image")
|
||||
image = Image.open(BytesIO(response.content))
|
||||
return image
|
||||
|
||||
def load_image_from_path(self, path):
|
||||
"""Loads an image from the given path"""
|
||||
return Image.open(path)
|
||||
|
||||
def split_image(
|
||||
self, image, h_splits: int = 2, v_splits: int = 2
|
||||
):
|
||||
"""Splits the given image into h_splits x v_splits parts"""
|
||||
width, height = image.size
|
||||
w_step, h_step = width // h_splits, height // v_splits
|
||||
slices = []
|
||||
|
||||
for i in range(v_splits):
|
||||
for j in range(h_splits):
|
||||
slice = image.crop(
|
||||
(
|
||||
j * w_step,
|
||||
i * h_step,
|
||||
(j + 1) * w_step,
|
||||
(i + 1) * h_step,
|
||||
)
|
||||
)
|
||||
slices.append(slice)
|
||||
return slices
|
||||
|
||||
def get_vectors(
|
||||
self,
|
||||
image,
|
||||
h_splits: int = 2,
|
||||
v_splits: int = 2,
|
||||
):
|
||||
"""Gets the vectors for the given image"""
|
||||
slices = self.split_image(image, h_splits, v_splits)
|
||||
vectors = []
|
||||
|
||||
for slice in slices:
|
||||
inputs = self.processor(
|
||||
text=self.query_text,
|
||||
images=slice,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
)
|
||||
outputs = self.model(**inputs)
|
||||
vectors.append(
|
||||
outputs.image_embeds.squeeze().detach().numpy()
|
||||
)
|
||||
return vectors
|
||||
|
||||
def run_from_url(
|
||||
self,
|
||||
url: str = "https://picsum.photos/800",
|
||||
h_splits: int = 2,
|
||||
v_splits: int = 2,
|
||||
):
|
||||
"""Runs the model on the image fetched from the given url"""
|
||||
image = self.fetch_image_from_url(url)
|
||||
return self.get_vectors(image, h_splits, v_splits)
|
||||
|
||||
def check_hard_chunking(self, quadrants):
|
||||
"""Check if the chunking is hard"""
|
||||
variances = []
|
||||
for quadrant in quadrants:
|
||||
edge_pixels = torch.cat(
|
||||
[
|
||||
quadrant[0, 1],
|
||||
quadrant[-1, :],
|
||||
]
|
||||
)
|
||||
variances.append(torch.var(edge_pixels).item())
|
||||
return variances
|
||||
|
||||
def embed_whole_image(self, image):
|
||||
"""Embed the entire image"""
|
||||
inputs = self.processor(
|
||||
image,
|
||||
return_tensors="pt",
|
||||
)
|
||||
with torch.no_grad():
|
||||
outputs = self.model(**inputs)
|
||||
return outputs.image_embeds.squeeze()
|
||||
|
||||
def apply_noise_reduction(self, image, kernel_size: int = 5):
|
||||
"""Implement an upscaling method to upscale the image and tiling issues"""
|
||||
blur = GaussianBlur(kernel_size)
|
||||
return blur(image)
|
||||
|
||||
def run_from_path(
|
||||
self, path: str = None, h_splits: int = 2, v_splits: int = 2
|
||||
):
|
||||
"""Runs the model on the image loaded from the given path"""
|
||||
image = self.load_image_from_path(path)
|
||||
return self.get_vectors(image, h_splits, v_splits)
|
||||
|
||||
def get_captions(self, image, candidate_captions):
|
||||
"""Get the best caption for the given image"""
|
||||
inputs_image = self.processor(
|
||||
images=image,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
inputs_text = self.processor(
|
||||
text=candidate_captions,
|
||||
images=inputs_image.pixel_values[
|
||||
0
|
||||
], # Fix the argument name
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
truncation=True,
|
||||
)
|
||||
|
||||
image_embeds = self.model(
|
||||
pixel_values=inputs_image.pixel_values[0]
|
||||
).image_embeds
|
||||
text_embeds = self.model(
|
||||
input_ids=inputs_text.input_ids,
|
||||
attention_mask=inputs_text.attention_mask,
|
||||
).text_embeds
|
||||
|
||||
# Calculate similarity between image and text
|
||||
similarities = (image_embeds @ text_embeds.T).squeeze(0)
|
||||
best_caption_index = similarities.argmax().item()
|
||||
|
||||
return candidate_captions[best_caption_index]
|
||||
|
||||
def get_and_concat_captions(
|
||||
self, image, candidate_captions, h_splits=2, v_splits=2
|
||||
):
|
||||
"""Get the best caption for the given image"""
|
||||
slices = self.split_image(image, h_splits, v_splits)
|
||||
captions = [
|
||||
self.get_captions(slice, candidate_captions)
|
||||
for slice in slices
|
||||
]
|
||||
concated_captions = "".join(captions)
|
||||
return concated_captions
|
||||
@ -1,130 +0,0 @@
|
||||
import torch
|
||||
from modelscope import AutoModelForCausalLM, AutoTokenizer
|
||||
from PIL import Image
|
||||
|
||||
from swarms.models.base_multimodal_model import BaseMultiModalModel
|
||||
|
||||
device_check = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
|
||||
class CogAgent(BaseMultiModalModel):
|
||||
"""CogAgent
|
||||
|
||||
Multi-modal conversational agent that can be used to chat with
|
||||
images and text. It is based on the CogAgent model from the
|
||||
ModelScope library.
|
||||
|
||||
Attributes:
|
||||
model_name (str): The name of the model to be used
|
||||
tokenizer_name (str): The name of the tokenizer to be used
|
||||
dtype (torch.bfloat16): The data type to be used
|
||||
low_cpu_mem_usage (bool): Whether to use low CPU memory
|
||||
load_in_4bit (bool): Whether to load in 4-bit
|
||||
trust_remote_code (bool): Whether to trust remote code
|
||||
device (str): The device to be used
|
||||
|
||||
Examples:
|
||||
>>> from swarms.models.cog_agent import CogAgent
|
||||
>>> cog_agent = CogAgent()
|
||||
>>> cog_agent.run("How are you?", "images/1.jpg")
|
||||
<s> I'm fine. How are you? </s>
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str = "ZhipuAI/cogagent-chat",
|
||||
tokenizer_name: str = "I-ModelScope/vicuna-7b-v1.5",
|
||||
dtype=torch.bfloat16,
|
||||
low_cpu_mem_usage: bool = True,
|
||||
load_in_4bit: bool = True,
|
||||
trust_remote_code: bool = True,
|
||||
device=device_check,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.model_name = model_name
|
||||
self.tokenizer_name = tokenizer_name
|
||||
self.dtype = dtype
|
||||
self.low_cpu_mem_usage = low_cpu_mem_usage
|
||||
self.load_in_4bit = load_in_4bit
|
||||
self.trust_remote_code = trust_remote_code
|
||||
self.device = device
|
||||
|
||||
self.model = (
|
||||
AutoModelForCausalLM.from_pretrained(
|
||||
self.model_name,
|
||||
torch_dtype=self.dtype,
|
||||
low_cpu_mem_usage=self.low_cpu_mem_usage,
|
||||
load_in_4bit=self.load_in_4bit,
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
.to(self.device)
|
||||
.eval()
|
||||
)
|
||||
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
self.tokenizer_name
|
||||
)
|
||||
|
||||
def run(self, task: str, img: str, *args, **kwargs):
|
||||
"""Run the model
|
||||
|
||||
Args:
|
||||
task (str): The task to be performed
|
||||
img (str): The image path
|
||||
|
||||
"""
|
||||
image = Image.open(img).convert("RGB")
|
||||
|
||||
input_by_model = self.model.build_conversation_input_ids(
|
||||
self.tokenizer,
|
||||
query=task,
|
||||
history=[],
|
||||
images=[image],
|
||||
)
|
||||
|
||||
inputs = {
|
||||
"input_ids": (
|
||||
input_by_model["input_ids"]
|
||||
.unsqueeze(0)
|
||||
.to(self.device)
|
||||
),
|
||||
"token_type_ids": (
|
||||
input_by_model["token_type_ids"]
|
||||
.unsqueeze(0)
|
||||
.to(self.device)
|
||||
),
|
||||
"attention_mask": (
|
||||
input_by_model["attention_mask"]
|
||||
.unsqueeze(0)
|
||||
.to(self.device)
|
||||
),
|
||||
"images": [
|
||||
[
|
||||
input_by_model["images"][0]
|
||||
.to(self.device)
|
||||
.to(self.dtype)
|
||||
]
|
||||
],
|
||||
}
|
||||
if (
|
||||
"cross_images" in input_by_model
|
||||
and input_by_model["cross_images"]
|
||||
):
|
||||
inputs["cross_images"] = [
|
||||
[
|
||||
input_by_model["cross_images"][0]
|
||||
.to(self.device)
|
||||
.to(self.dtype)
|
||||
]
|
||||
]
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = self.model(**inputs, **kwargs)
|
||||
outputs = outputs[:, inputs["input_ids"].shape[1] :]
|
||||
response = self.decode(outputs[0])
|
||||
response = response.split("</s>")[0]
|
||||
print(response)
|
||||
@ -1,106 +0,0 @@
|
||||
import requests
|
||||
from tenacity import retry, stop_after_attempt, wait_fixed
|
||||
|
||||
|
||||
class Gigabind:
|
||||
"""Gigabind API.
|
||||
|
||||
Args:
|
||||
host (str, optional): host. Defaults to None.
|
||||
proxy_url (str, optional): proxy_url. Defaults to None.
|
||||
port (int, optional): port. Defaults to 8000.
|
||||
endpoint (str, optional): endpoint. Defaults to "embeddings".
|
||||
|
||||
Examples:
|
||||
>>> from swarms.models.gigabind import Gigabind
|
||||
>>> api = Gigabind(host="localhost", port=8000, endpoint="embeddings")
|
||||
>>> response = api.run(text="Hello, world!", vision="image.jpg")
|
||||
>>> print(response)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = None,
|
||||
proxy_url: str = None,
|
||||
port: int = 8000,
|
||||
endpoint: str = "embeddings",
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.host = host
|
||||
self.proxy_url = proxy_url
|
||||
self.port = port
|
||||
self.endpoint = endpoint
|
||||
|
||||
# Set the URL to the API
|
||||
if self.proxy_url is not None:
|
||||
self.url = f"{self.proxy_url}"
|
||||
else:
|
||||
self.url = f"http://{host}:{port}/{endpoint}"
|
||||
|
||||
@retry(stop=stop_after_attempt(3), wait=wait_fixed(2))
|
||||
def run(
|
||||
self,
|
||||
text: str = None,
|
||||
vision: str = None,
|
||||
audio: str = None,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
"""Run the Gigabind API.
|
||||
|
||||
Args:
|
||||
text (str, optional): text. Defaults to None.
|
||||
vision (str, optional): images. Defaults to None.
|
||||
audio (str, optional): audio file paths. Defaults to None.
|
||||
|
||||
Raises:
|
||||
ValueError: At least one of text, vision or audio must be provided
|
||||
|
||||
Returns:
|
||||
embeddings: embeddings
|
||||
"""
|
||||
try:
|
||||
# Prepare the data to send to the API
|
||||
data = {}
|
||||
if text is not None:
|
||||
data["text"] = text
|
||||
if vision is not None:
|
||||
data["vision"] = vision
|
||||
if audio is not None:
|
||||
data["audio"] = audio
|
||||
else:
|
||||
raise ValueError(
|
||||
"At least one of text, vision or audio must be"
|
||||
" provided"
|
||||
)
|
||||
|
||||
# Send a POST request to the API and return the response
|
||||
response = requests.post(
|
||||
self.url, json=data, *args, **kwargs
|
||||
)
|
||||
return response.json()
|
||||
except Exception as error:
|
||||
print(f"Gigabind API error: {error}")
|
||||
return None
|
||||
|
||||
def generate_summary(self, text: str = None, *args, **kwargs):
|
||||
# Prepare the data to send to the API
|
||||
data = {}
|
||||
if text is not None:
|
||||
data["text"] = text
|
||||
else:
|
||||
raise ValueError(
|
||||
"At least one of text, vision or audio must be"
|
||||
" provided"
|
||||
)
|
||||
|
||||
# Send a POST request to the API and return the response
|
||||
response = requests.post(self.url, json=data, *args, **kwargs)
|
||||
return response.json()
|
||||
|
||||
|
||||
# api = Gigabind(host="localhost", port=8000, endpoint="embeddings")
|
||||
# response = api.run(text="Hello, world!", vision="image.jpg")
|
||||
# print(response)
|
||||
Loading…
Reference in new issue