diff --git a/demos/accountant_team/accountant_team.py b/demos/accountant_team/accountant_team.py index 459da830..18bda007 100644 --- a/demos/accountant_team/accountant_team.py +++ b/demos/accountant_team/accountant_team.py @@ -1,7 +1,6 @@ # !pip install --upgrade swarms==2.0.6 - from swarms.models import OpenAIChat from swarms.models.nougat import Nougat from swarms.structs import Flow @@ -11,7 +10,10 @@ from swarms.structs.sequential_workflow import SequentialWorkflow IMAGE_OF_FINANCIAL_DOC_URL = "bank_statement_2.jpg" # Example usage -api_key = "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" # Your actual API key here +api_key = ( + "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" # Your actual API key here +) + # Initialize the OCR model def ocr_model(img: str): @@ -19,12 +21,14 @@ def ocr_model(img: str): analyze_finance_docs = ocr(img) return str(analyze_finance_docs) + # Initialize the language flow llm = OpenAIChat( openai_api_key=api_key, temperature=0.5, ) + # Create a prompt for the language model def summary_agent_prompt(analyzed_doc: str): analyzed_doc = ocr_model(img=analyzed_doc) @@ -36,6 +40,7 @@ def summary_agent_prompt(analyzed_doc: str): {analyzed_doc} """ + # Initialize the Flow with the language flow flow1 = Flow(llm=llm, max_loops=1, dashboard=False) @@ -49,11 +54,14 @@ workflow = SequentialWorkflow(max_loops=1) workflow.add(summary_agent_prompt(IMAGE_OF_FINANCIAL_DOC_URL), flow1) # Suppose the next task takes the output of the first task as input -workflow.add("Provide an actionable step by step plan on how to cut costs from the analyzed financial document.", flow2) +workflow.add( + "Provide an actionable step by step plan on how to cut costs from the analyzed financial document.", + flow2, +) # Run the workflow workflow.run() # Output the results for task in workflow.tasks: - print(f"Task: {task.description}, Result: {task.result}") \ No newline at end of file + print(f"Task: {task.description}, Result: {task.result}") diff --git a/playground/models/anthropic_example.py b/playground/models/anthropic_example.py index 695dfe62..3354d0cc 100644 --- a/playground/models/anthropic_example.py +++ b/playground/models/anthropic_example.py @@ -1,11 +1,9 @@ from swarms.models.anthropic import Anthropic -model = Anthropic( - anthropic_api_key="" -) +model = Anthropic(anthropic_api_key="") task = "Say hello to" -print(model(task)) \ No newline at end of file +print(model(task)) diff --git a/swarms/models/anthropic.py b/swarms/models/anthropic.py index 24c3c126..c8ba2329 100644 --- a/swarms/models/anthropic.py +++ b/swarms/models/anthropic.py @@ -523,4 +523,4 @@ class Anthropic(LLM, _AnthropicCommon): """Calculate number of tokens.""" if not self.count_tokens: raise NameError("Please ensure the anthropic package is loaded") - return self.count_tokens(text) \ No newline at end of file + return self.count_tokens(text) diff --git a/swarms/models/fast_vit_classes.json b/swarms/models/fast_vit_classes.json new file mode 100644 index 00000000..57434253 --- /dev/null +++ b/swarms/models/fast_vit_classes.json @@ -0,0 +1,1000 @@ +["tench", +"goldfish", +"great white shark", +"tiger shark", +"hammerhead shark", +"electric ray", +"stingray", +"cock", +"hen", +"ostrich", +"brambling", +"goldfinch", +"house finch", +"junco", +"indigo bunting", +"American robin", +"bulbul", +"jay", +"magpie", +"chickadee", +"American dipper", +"kite", +"bald eagle", +"vulture", +"great grey owl", +"fire salamander", +"smooth newt", +"newt", +"spotted salamander", +"axolotl", +"American bullfrog", +"tree frog", +"tailed frog", +"loggerhead sea turtle", +"leatherback sea turtle", +"mud turtle", +"terrapin", +"box turtle", +"banded gecko", +"green iguana", +"Carolina anole", +"desert grassland whiptail lizard", +"agama", +"frilled-necked lizard", +"alligator lizard", +"Gila monster", +"European green lizard", +"chameleon", +"Komodo dragon", +"Nile crocodile", +"American alligator", +"triceratops", +"worm snake", +"ring-necked snake", +"eastern hog-nosed snake", +"smooth green snake", +"kingsnake", +"garter snake", +"water snake", +"vine snake", +"night snake", +"boa constrictor", +"African rock python", +"Indian cobra", +"green mamba", +"sea snake", +"Saharan horned viper", +"eastern diamondback rattlesnake", +"sidewinder", +"trilobite", +"harvestman", +"scorpion", +"yellow garden spider", +"barn spider", +"European garden spider", +"southern black widow", +"tarantula", +"wolf spider", +"tick", +"centipede", +"black grouse", +"ptarmigan", +"ruffed grouse", +"prairie grouse", +"peacock", +"quail", +"partridge", +"grey parrot", +"macaw", +"sulphur-crested cockatoo", +"lorikeet", +"coucal", +"bee eater", +"hornbill", +"hummingbird", +"jacamar", +"toucan", +"duck", +"red-breasted merganser", +"goose", +"black swan", +"tusker", +"echidna", +"platypus", +"wallaby", +"koala", +"wombat", +"jellyfish", +"sea anemone", +"brain coral", +"flatworm", +"nematode", +"conch", +"snail", +"slug", +"sea slug", +"chiton", +"chambered nautilus", +"Dungeness crab", +"rock crab", +"fiddler crab", +"red king crab", +"American lobster", +"spiny lobster", +"crayfish", +"hermit crab", +"isopod", +"white stork", +"black stork", +"spoonbill", +"flamingo", +"little blue heron", +"great egret", +"bittern", +"crane (bird)", +"limpkin", +"common gallinule", +"American coot", +"bustard", +"ruddy turnstone", +"dunlin", +"common redshank", +"dowitcher", +"oystercatcher", +"pelican", +"king penguin", +"albatross", +"grey whale", +"killer whale", +"dugong", +"sea lion", +"Chihuahua", +"Japanese Chin", +"Maltese", +"Pekingese", +"Shih Tzu", +"King Charles Spaniel", +"Papillon", +"toy terrier", +"Rhodesian Ridgeback", +"Afghan Hound", +"Basset Hound", +"Beagle", +"Bloodhound", +"Bluetick Coonhound", +"Black and Tan Coonhound", +"Treeing Walker Coonhound", +"English foxhound", +"Redbone Coonhound", +"borzoi", +"Irish Wolfhound", +"Italian Greyhound", +"Whippet", +"Ibizan Hound", +"Norwegian Elkhound", +"Otterhound", +"Saluki", +"Scottish Deerhound", +"Weimaraner", +"Staffordshire Bull Terrier", +"American Staffordshire Terrier", +"Bedlington Terrier", +"Border Terrier", +"Kerry Blue Terrier", +"Irish Terrier", +"Norfolk Terrier", +"Norwich Terrier", +"Yorkshire Terrier", +"Wire Fox Terrier", +"Lakeland Terrier", +"Sealyham Terrier", +"Airedale Terrier", +"Cairn Terrier", +"Australian Terrier", +"Dandie Dinmont Terrier", +"Boston Terrier", +"Miniature Schnauzer", +"Giant Schnauzer", +"Standard Schnauzer", +"Scottish Terrier", +"Tibetan Terrier", +"Australian Silky Terrier", +"Soft-coated Wheaten Terrier", +"West Highland White Terrier", +"Lhasa Apso", +"Flat-Coated Retriever", +"Curly-coated Retriever", +"Golden Retriever", +"Labrador Retriever", +"Chesapeake Bay Retriever", +"German Shorthaired Pointer", +"Vizsla", +"English Setter", +"Irish Setter", +"Gordon Setter", +"Brittany Spaniel", +"Clumber Spaniel", +"English Springer Spaniel", +"Welsh Springer Spaniel", +"Cocker Spaniels", +"Sussex Spaniel", +"Irish Water Spaniel", +"Kuvasz", +"Schipperke", +"Groenendael", +"Malinois", +"Briard", +"Australian Kelpie", +"Komondor", +"Old English Sheepdog", +"Shetland Sheepdog", +"collie", +"Border Collie", +"Bouvier des Flandres", +"Rottweiler", +"German Shepherd Dog", +"Dobermann", +"Miniature Pinscher", +"Greater Swiss Mountain Dog", +"Bernese Mountain Dog", +"Appenzeller Sennenhund", +"Entlebucher Sennenhund", +"Boxer", +"Bullmastiff", +"Tibetan Mastiff", +"French Bulldog", +"Great Dane", +"St. Bernard", +"husky", +"Alaskan Malamute", +"Siberian Husky", +"Dalmatian", +"Affenpinscher", +"Basenji", +"pug", +"Leonberger", +"Newfoundland", +"Pyrenean Mountain Dog", +"Samoyed", +"Pomeranian", +"Chow Chow", +"Keeshond", +"Griffon Bruxellois", +"Pembroke Welsh Corgi", +"Cardigan Welsh Corgi", +"Toy Poodle", +"Miniature Poodle", +"Standard Poodle", +"Mexican hairless dog", +"grey wolf", +"Alaskan tundra wolf", +"red wolf", +"coyote", +"dingo", +"dhole", +"African wild dog", +"hyena", +"red fox", +"kit fox", +"Arctic fox", +"grey fox", +"tabby cat", +"tiger cat", +"Persian cat", +"Siamese cat", +"Egyptian Mau", +"cougar", +"lynx", +"leopard", +"snow leopard", +"jaguar", +"lion", +"tiger", +"cheetah", +"brown bear", +"American black bear", +"polar bear", +"sloth bear", +"mongoose", +"meerkat", +"tiger beetle", +"ladybug", +"ground beetle", +"longhorn beetle", +"leaf beetle", +"dung beetle", +"rhinoceros beetle", +"weevil", +"fly", +"bee", +"ant", +"grasshopper", +"cricket", +"stick insect", +"cockroach", +"mantis", +"cicada", +"leafhopper", +"lacewing", +"dragonfly", +"damselfly", +"red admiral", +"ringlet", +"monarch butterfly", +"small white", +"sulphur butterfly", +"gossamer-winged butterfly", +"starfish", +"sea urchin", +"sea cucumber", +"cottontail rabbit", +"hare", +"Angora rabbit", +"hamster", +"porcupine", +"fox squirrel", +"marmot", +"beaver", +"guinea pig", +"common sorrel", +"zebra", +"pig", +"wild boar", +"warthog", +"hippopotamus", +"ox", +"water buffalo", +"bison", +"ram", +"bighorn sheep", +"Alpine ibex", +"hartebeest", +"impala", +"gazelle", +"dromedary", +"llama", +"weasel", +"mink", +"European polecat", +"black-footed ferret", +"otter", +"skunk", +"badger", +"armadillo", +"three-toed sloth", +"orangutan", +"gorilla", +"chimpanzee", +"gibbon", +"siamang", +"guenon", +"patas monkey", +"baboon", +"macaque", +"langur", +"black-and-white colobus", +"proboscis monkey", +"marmoset", +"white-headed capuchin", +"howler monkey", +"titi", +"Geoffroy's spider monkey", +"common squirrel monkey", +"ring-tailed lemur", +"indri", +"Asian elephant", +"African bush elephant", +"red panda", +"giant panda", +"snoek", +"eel", +"coho salmon", +"rock beauty", +"clownfish", +"sturgeon", +"garfish", +"lionfish", +"pufferfish", +"abacus", +"abaya", +"academic gown", +"accordion", +"acoustic guitar", +"aircraft carrier", +"airliner", +"airship", +"altar", +"ambulance", +"amphibious vehicle", +"analog clock", +"apiary", +"apron", +"waste container", +"assault rifle", +"backpack", +"bakery", +"balance beam", +"balloon", +"ballpoint pen", +"Band-Aid", +"banjo", +"baluster", +"barbell", +"barber chair", +"barbershop", +"barn", +"barometer", +"barrel", +"wheelbarrow", +"baseball", +"basketball", +"bassinet", +"bassoon", +"swimming cap", +"bath towel", +"bathtub", +"station wagon", +"lighthouse", +"beaker", +"military cap", +"beer bottle", +"beer glass", +"bell-cot", +"bib", +"tandem bicycle", +"bikini", +"ring binder", +"binoculars", +"birdhouse", +"boathouse", +"bobsleigh", +"bolo tie", +"poke bonnet", +"bookcase", +"bookstore", +"bottle cap", +"bow", +"bow tie", +"brass", +"bra", +"breakwater", +"breastplate", +"broom", +"bucket", +"buckle", +"bulletproof vest", +"high-speed train", +"butcher shop", +"taxicab", +"cauldron", +"candle", +"cannon", +"canoe", +"can opener", +"cardigan", +"car mirror", +"carousel", +"tool kit", +"carton", +"car wheel", +"automated teller machine", +"cassette", +"cassette player", +"castle", +"catamaran", +"CD player", +"cello", +"mobile phone", +"chain", +"chain-link fence", +"chain mail", +"chainsaw", +"chest", +"chiffonier", +"chime", +"china cabinet", +"Christmas stocking", +"church", +"movie theater", +"cleaver", +"cliff dwelling", +"cloak", +"clogs", +"cocktail shaker", +"coffee mug", +"coffeemaker", +"coil", +"combination lock", +"computer keyboard", +"confectionery store", +"container ship", +"convertible", +"corkscrew", +"cornet", +"cowboy boot", +"cowboy hat", +"cradle", +"crane (machine)", +"crash helmet", +"crate", +"infant bed", +"Crock Pot", +"croquet ball", +"crutch", +"cuirass", +"dam", +"desk", +"desktop computer", +"rotary dial telephone", +"diaper", +"digital clock", +"digital watch", +"dining table", +"dishcloth", +"dishwasher", +"disc brake", +"dock", +"dog sled", +"dome", +"doormat", +"drilling rig", +"drum", +"drumstick", +"dumbbell", +"Dutch oven", +"electric fan", +"electric guitar", +"electric locomotive", +"entertainment center", +"envelope", +"espresso machine", +"face powder", +"feather boa", +"filing cabinet", +"fireboat", +"fire engine", +"fire screen sheet", +"flagpole", +"flute", +"folding chair", +"football helmet", +"forklift", +"fountain", +"fountain pen", +"four-poster bed", +"freight car", +"French horn", +"frying pan", +"fur coat", +"garbage truck", +"gas mask", +"gas pump", +"goblet", +"go-kart", +"golf ball", +"golf cart", +"gondola", +"gong", +"gown", +"grand piano", +"greenhouse", +"grille", +"grocery store", +"guillotine", +"barrette", +"hair spray", +"half-track", +"hammer", +"hamper", +"hair dryer", +"hand-held computer", +"handkerchief", +"hard disk drive", +"harmonica", +"harp", +"harvester", +"hatchet", +"holster", +"home theater", +"honeycomb", +"hook", +"hoop skirt", +"horizontal bar", +"horse-drawn vehicle", +"hourglass", +"iPod", +"clothes iron", +"jack-o'-lantern", +"jeans", +"jeep", +"T-shirt", +"jigsaw puzzle", +"pulled rickshaw", +"joystick", +"kimono", +"knee pad", +"knot", +"lab coat", +"ladle", +"lampshade", +"laptop computer", +"lawn mower", +"lens cap", +"paper knife", +"library", +"lifeboat", +"lighter", +"limousine", +"ocean liner", +"lipstick", +"slip-on shoe", +"lotion", +"speaker", +"loupe", +"sawmill", +"magnetic compass", +"mail bag", +"mailbox", +"tights", +"tank suit", +"manhole cover", +"maraca", +"marimba", +"mask", +"match", +"maypole", +"maze", +"measuring cup", +"medicine chest", +"megalith", +"microphone", +"microwave oven", +"military uniform", +"milk can", +"minibus", +"miniskirt", +"minivan", +"missile", +"mitten", +"mixing bowl", +"mobile home", +"Model T", +"modem", +"monastery", +"monitor", +"moped", +"mortar", +"square academic cap", +"mosque", +"mosquito net", +"scooter", +"mountain bike", +"tent", +"computer mouse", +"mousetrap", +"moving van", +"muzzle", +"nail", +"neck brace", +"necklace", +"nipple", +"notebook computer", +"obelisk", +"oboe", +"ocarina", +"odometer", +"oil filter", +"organ", +"oscilloscope", +"overskirt", +"bullock cart", +"oxygen mask", +"packet", +"paddle", +"paddle wheel", +"padlock", +"paintbrush", +"pajamas", +"palace", +"pan flute", +"paper towel", +"parachute", +"parallel bars", +"park bench", +"parking meter", +"passenger car", +"patio", +"payphone", +"pedestal", +"pencil case", +"pencil sharpener", +"perfume", +"Petri dish", +"photocopier", +"plectrum", +"Pickelhaube", +"picket fence", +"pickup truck", +"pier", +"piggy bank", +"pill bottle", +"pillow", +"ping-pong ball", +"pinwheel", +"pirate ship", +"pitcher", +"hand plane", +"planetarium", +"plastic bag", +"plate rack", +"plow", +"plunger", +"Polaroid camera", +"pole", +"police van", +"poncho", +"billiard table", +"soda bottle", +"pot", +"potter's wheel", +"power drill", +"prayer rug", +"printer", +"prison", +"projectile", +"projector", +"hockey puck", +"punching bag", +"purse", +"quill", +"quilt", +"race car", +"racket", +"radiator", +"radio", +"radio telescope", +"rain barrel", +"recreational vehicle", +"reel", +"reflex camera", +"refrigerator", +"remote control", +"restaurant", +"revolver", +"rifle", +"rocking chair", +"rotisserie", +"eraser", +"rugby ball", +"ruler", +"running shoe", +"safe", +"safety pin", +"salt shaker", +"sandal", +"sarong", +"saxophone", +"scabbard", +"weighing scale", +"school bus", +"schooner", +"scoreboard", +"CRT screen", +"screw", +"screwdriver", +"seat belt", +"sewing machine", +"shield", +"shoe store", +"shoji", +"shopping basket", +"shopping cart", +"shovel", +"shower cap", +"shower curtain", +"ski", +"ski mask", +"sleeping bag", +"slide rule", +"sliding door", +"slot machine", +"snorkel", +"snowmobile", +"snowplow", +"soap dispenser", +"soccer ball", +"sock", +"solar thermal collector", +"sombrero", +"soup bowl", +"space bar", +"space heater", +"space shuttle", +"spatula", +"motorboat", +"spider web", +"spindle", +"sports car", +"spotlight", +"stage", +"steam locomotive", +"through arch bridge", +"steel drum", +"stethoscope", +"scarf", +"stone wall", +"stopwatch", +"stove", +"strainer", +"tram", +"stretcher", +"couch", +"stupa", +"submarine", +"suit", +"sundial", +"sunglass", +"sunglasses", +"sunscreen", +"suspension bridge", +"mop", +"sweatshirt", +"swimsuit", +"swing", +"switch", +"syringe", +"table lamp", +"tank", +"tape player", +"teapot", +"teddy bear", +"television", +"tennis ball", +"thatched roof", +"front curtain", +"thimble", +"threshing machine", +"throne", +"tile roof", +"toaster", +"tobacco shop", +"toilet seat", +"torch", +"totem pole", +"tow truck", +"toy store", +"tractor", +"semi-trailer truck", +"tray", +"trench coat", +"tricycle", +"trimaran", +"tripod", +"triumphal arch", +"trolleybus", +"trombone", +"tub", +"turnstile", +"typewriter keyboard", +"umbrella", +"unicycle", +"upright piano", +"vacuum cleaner", +"vase", +"vault", +"velvet", +"vending machine", +"vestment", +"viaduct", +"violin", +"volleyball", +"waffle iron", +"wall clock", +"wallet", +"wardrobe", +"military aircraft", +"sink", +"washing machine", +"water bottle", +"water jug", +"water tower", +"whiskey jug", +"whistle", +"wig", +"window screen", +"window shade", +"Windsor tie", +"wine bottle", +"wing", +"wok", +"wooden spoon", +"wool", +"split-rail fence", +"shipwreck", +"yawl", +"yurt", +"website", +"comic book", +"crossword", +"traffic sign", +"traffic light", +"dust jacket", +"menu", +"plate", +"guacamole", +"consomme", +"hot pot", +"trifle", +"ice cream", +"ice pop", +"baguette", +"bagel", +"pretzel", +"cheeseburger", +"hot dog", +"mashed potato", +"cabbage", +"broccoli", +"cauliflower", +"zucchini", +"spaghetti squash", +"acorn squash", +"butternut squash", +"cucumber", +"artichoke", +"bell pepper", +"cardoon", +"mushroom", +"Granny Smith", +"strawberry", +"orange", +"lemon", +"fig", +"pineapple", +"banana", +"jackfruit", +"custard apple", +"pomegranate", +"hay", +"carbonara", +"chocolate syrup", +"dough", +"meatloaf", +"pizza", +"pot pie", +"burrito", +"red wine", +"espresso", +"cup", +"eggnog", +"alp", +"bubble", +"cliff", +"coral reef", +"geyser", +"lakeshore", +"promontory", +"shoal", +"seashore", +"valley", +"volcano", +"baseball player", +"bridegroom", +"scuba diver", +"rapeseed", +"daisy", +"yellow lady's slipper", +"corn", +"acorn", +"rose hip", +"horse chestnut seed", +"coral fungus", +"agaric", +"gyromitra", +"stinkhorn mushroom", +"earth star", +"hen-of-the-woods", +"bolete", +"ear of corn", +"toilet paper"] \ No newline at end of file diff --git a/swarms/models/fastvit.py b/swarms/models/fastvit.py new file mode 100644 index 00000000..a9a6abf7 --- /dev/null +++ b/swarms/models/fastvit.py @@ -0,0 +1,80 @@ +import json +import os +from typing import List + +import numpy as np +import timm +import torch +from PIL import Image +from pydantic import BaseModel, StrictFloat, StrictInt, validator + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +# Load the classes for image classification +with open(os.path.join(os.path.dirname(__file__), "fast_vit_classes.json")) as f: + FASTVIT_IMAGENET_1K_CLASSES = json.load(f) + + +class ClassificationResult(BaseModel): + class_id: List[StrictInt] + confidence: List[StrictFloat] + + @validator("class_id", "confidence", pre=True, each_item=True) + def check_list_contents(cls, v): + assert isinstance(v, int) or isinstance(v, float), "must be integer or float" + return v + + +class FastViT: + """ + FastViT model for image classification + + Args: + img (str): path to the input image + confidence_threshold (float): confidence threshold for the model's predictions + + Returns: + ClassificationResult: a pydantic BaseModel containing the class ids and confidences of the model's predictions + + + Example: + >>> fastvit = FastViT() + >>> result = fastvit(img="path_to_image.jpg", confidence_threshold=0.5) + + + To use, create a json file called: fast_vit_classes.json + + """ + def __init__(self): + self.model = timm.create_model( + "hf_hub:timm/fastvit_s12.apple_in1k", pretrained=True + ).to(DEVICE) + data_config = timm.data.resolve_model_data_config(self.model) + self.transforms = timm.data.create_transform(**data_config, is_training=False) + self.model.eval() + + def __call__( + self, img: str, confidence_threshold: float = 0.5 + ) -> ClassificationResult: + """classifies the input image and returns the top k classes and their probabilities""" + img = Image.open(img).convert("RGB") + img_tensor = self.transforms(img).unsqueeze(0).to(DEVICE) + with torch.no_grad(): + output = self.model(img_tensor) + probabilities = torch.nn.functional.softmax(output, dim=1) + + # Get top k classes and their probabilities + top_probs, top_classes = torch.topk( + probabilities, k=FASTVIT_IMAGENET_1K_CLASSES + ) + + # Filter by confidence threshold + mask = top_probs > confidence_threshold + top_probs, top_classes = top_probs[mask], top_classes[mask] + + # Convert to Python lists and map class indices to labels if needed + top_probs = top_probs.cpu().numpy().tolist() + top_classes = top_classes.cpu().numpy().tolist() + # top_class_labels = [FASTVIT_IMAGENET_1K_CLASSES[i] for i in top_classes] # Uncomment if class labels are needed + + return ClassificationResult(class_id=top_classes, confidence=top_probs) diff --git a/swarms/models/kosmos2.py b/swarms/models/kosmos2.py new file mode 100644 index 00000000..12d5638a --- /dev/null +++ b/swarms/models/kosmos2.py @@ -0,0 +1,100 @@ +from typing import List, Tuple + +import numpy as np +from PIL import Image +from pydantic import BaseModel, root_validator, validator +from transformers import AutoModelForVision2Seq, AutoProcessor + + +# Assuming the Detections class represents the output of the model prediction +class Detections(BaseModel): + xyxy: List[Tuple[float, float, float, float]] + class_id: List[int] + confidence: List[float] + + @root_validator + def check_length(cls, values): + assert ( + len(values.get("xyxy")) + == len(values.get("class_id")) + == len(values.get("confidence")) + ), "All fields must have the same length." + return values + + @validator("xyxy", "class_id", "confidence", pre=True, each_item=True) + def check_not_empty(cls, v): + if isinstance(v, list) and len(v) == 0: + raise ValueError("List must not be empty") + return v + + @classmethod + def empty(cls): + return cls(xyxy=[], class_id=[], confidence=[]) + + +class Kosmos2(BaseModel): + model: AutoModelForVision2Seq + processor: AutoProcessor + + @classmethod + def initialize(cls): + model = AutoModelForVision2Seq.from_pretrained( + "ydshieh/kosmos-2-patch14-224", trust_remote_code=True + ) + processor = AutoProcessor.from_pretrained( + "ydshieh/kosmos-2-patch14-224", trust_remote_code=True + ) + return cls(model=model, processor=processor) + + def __call__(self, img: str) -> Detections: + image = Image.open(img) + prompt = "An image of" + + inputs = self.processor(text=prompt, images=image, return_tensors="pt") + outputs = self.model.generate(**inputs, use_cache=True, max_new_tokens=64) + + generated_text = self.processor.batch_decode(outputs, skip_special_tokens=True)[ + 0 + ] + + # The actual processing of generated_text to entities would go here + # For the purpose of this example, assume a mock function 'extract_entities' exists: + entities = self.extract_entities(generated_text) + + # Convert entities to detections format + detections = self.process_entities_to_detections(entities, image) + return detections + + def extract_entities( + self, text: str + ) -> List[Tuple[str, Tuple[float, float, float, float]]]: + # Placeholder function for entity extraction + # This should be replaced with the actual method of extracting entities + return [] + + def process_entities_to_detections( + self, + entities: List[Tuple[str, Tuple[float, float, float, float]]], + image: Image.Image, + ) -> Detections: + if not entities: + return Detections.empty() + + class_ids = [0] * len(entities) # Replace with actual class ID extraction logic + xyxys = [ + ( + e[1][0] * image.width, + e[1][1] * image.height, + e[1][2] * image.width, + e[1][3] * image.height, + ) + for e in entities + ] + confidences = [1.0] * len(entities) # Placeholder confidence + + return Detections(xyxy=xyxys, class_id=class_ids, confidence=confidences) + + +# Usage: +# kosmos2 = Kosmos2.initialize() +# detections = kosmos2(img="path_to_image.jpg") diff --git a/swarms/structs/sequential_workflow.py b/swarms/structs/sequential_workflow.py index 251447ba..a7e1cd63 100644 --- a/swarms/structs/sequential_workflow.py +++ b/swarms/structs/sequential_workflow.py @@ -265,7 +265,7 @@ class SequentialWorkflow: attrs=["bold", "underline"], ) ) - + def workflow_shutdown(self, **kwargs) -> None: print( colored( @@ -286,11 +286,13 @@ class SequentialWorkflow: ) ) - task = Task(description=task, flow=kwargs["flow"], args=list(kwargs["args"]), kwargs=kwargs["kwargs"]) + task = Task( + description=task, + flow=kwargs["flow"], + args=list(kwargs["args"]), + kwargs=kwargs["kwargs"], + ) self.tasks.append(task) - - - def load_workflow_state(self, filepath: str = None, **kwargs) -> None: """