From d4b12e0beacc1525820ec3eca51865642d4cd117 Mon Sep 17 00:00:00 2001 From: Kye Date: Wed, 8 Nov 2023 14:52:53 -0500 Subject: [PATCH] fuyu + zephyr fixes Former-commit-id: 9aa40842a63ca72b8229c9481b2dc6f24fec4abc --- demos/accountant_team/accountant_team.py | 43 +++++---------- playground/models/anthropic_example.py | 2 +- pyproject.toml | 2 +- swarms/models/fastvit.py | 3 +- swarms/models/fuyu.py | 7 +-- swarms/models/timm.py | 67 ++++++++++++++++++++++++ swarms/models/zephyr.py | 53 +++++++++++++++++-- 7 files changed, 136 insertions(+), 41 deletions(-) create mode 100644 swarms/models/timm.py diff --git a/demos/accountant_team/accountant_team.py b/demos/accountant_team/accountant_team.py index 18bda007..06f89684 100644 --- a/demos/accountant_team/accountant_team.py +++ b/demos/accountant_team/accountant_team.py @@ -1,43 +1,36 @@ # !pip install --upgrade swarms==2.0.6 - -from swarms.models import OpenAIChat +from swarms.models import BioGPT from swarms.models.nougat import Nougat from swarms.structs import Flow -from swarms.structs.sequential_workflow import SequentialWorkflow # # URL of the image of the financial document IMAGE_OF_FINANCIAL_DOC_URL = "bank_statement_2.jpg" # Example usage -api_key = ( - "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" # Your actual API key here -) - +api_key = "" # Your actual API key here # Initialize the OCR model -def ocr_model(img: str): - ocr = Nougat() - analyze_finance_docs = ocr(img) - return str(analyze_finance_docs) # Initialize the language flow -llm = OpenAIChat( - openai_api_key=api_key, - temperature=0.5, -) +llm = BioGPT() # Create a prompt for the language model def summary_agent_prompt(analyzed_doc: str): - analyzed_doc = ocr_model(img=analyzed_doc) + model = Nougat( + max_new_tokens=5000, + ) + + out = model(analyzed_doc) + return f""" Generate an actionable summary of this financial document, provide bulletpoints: Here is the Analyzed Document: --- - {analyzed_doc} + {out} """ @@ -47,21 +40,11 @@ flow1 = Flow(llm=llm, max_loops=1, dashboard=False) # Create another Flow for a different task flow2 = Flow(llm=llm, max_loops=1, dashboard=False) -# Create the workflow -workflow = SequentialWorkflow(max_loops=1) # Add tasks to the workflow -workflow.add(summary_agent_prompt(IMAGE_OF_FINANCIAL_DOC_URL), flow1) +summary_agent = flow1.run(summary_agent_prompt(IMAGE_OF_FINANCIAL_DOC_URL)) # 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, +out = flow2.run( + f"Provide an actionable step by step plan on how to cut costs from the analyzed financial document. {summary_agent}" ) - -# Run the workflow -workflow.run() - -# Output the results -for task in workflow.tasks: - print(f"Task: {task.description}, Result: {task.result}") diff --git a/playground/models/anthropic_example.py b/playground/models/anthropic_example.py index 3354d0cc..940892ca 100644 --- a/playground/models/anthropic_example.py +++ b/playground/models/anthropic_example.py @@ -4,6 +4,6 @@ from swarms.models.anthropic import Anthropic model = Anthropic(anthropic_api_key="") -task = "Say hello to" +task = "What is quantum field theory? What are 3 books on the field?" print(model(task)) diff --git a/pyproject.toml b/pyproject.toml index 3709ebc0..75f0e7ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api" [tool.poetry] name = "swarms" -version = "2.0.8" +version = "2.1.0" description = "Swarms - Pytorch" license = "MIT" authors = ["Kye Gomez "] diff --git a/swarms/models/fastvit.py b/swarms/models/fastvit.py index a9a6abf7..a2d6bc0a 100644 --- a/swarms/models/fastvit.py +++ b/swarms/models/fastvit.py @@ -43,8 +43,9 @@ class FastViT: 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 diff --git a/swarms/models/fuyu.py b/swarms/models/fuyu.py index 0fd1fd85..edecc19c 100644 --- a/swarms/models/fuyu.py +++ b/swarms/models/fuyu.py @@ -1,8 +1,8 @@ """Fuyu model by Kye""" from transformers import ( + FuyuProcessor, FuyuForCausalLM, AutoTokenizer, - FuyuProcessor, FuyuImageProcessor, ) from PIL import Image @@ -50,9 +50,9 @@ class Fuyu: pretrained_path, device_map=device_map ) - def __call__(self, text: str, img_path: str): + def __call__(self, text: str, img: str): """Call the model with text and img paths""" - image_pil = Image.open(img_path) + image_pil = Image.open(img) model_inputs = self.processor( text=text, images=[image_pil], device=self.device_map ) @@ -62,3 +62,4 @@ class Fuyu: output = self.model.generate(**model_inputs, max_new_tokens=self.max_new_tokens) text = self.processor.batch_decode(output[:, -7:], skip_special_tokens=True) + return print(str(text)) diff --git a/swarms/models/timm.py b/swarms/models/timm.py new file mode 100644 index 00000000..5d9b965a --- /dev/null +++ b/swarms/models/timm.py @@ -0,0 +1,67 @@ +from typing import List + +import timm +import torch +from pydantic import BaseModel, conlist + + +class TimmModelInfo(BaseModel): + model_name: str + pretrained: bool + in_chans: int + + class Config: + # Use strict typing for all fields + strict = True + + +class TimmModel: + """ + + # Usage + model_handler = TimmModelHandler() + model_info = TimmModelInfo(model_name='resnet34', pretrained=True, in_chans=1) + input_tensor = torch.randn(1, 1, 224, 224) + output_shape = model_handler(model_info=model_info, input_tensor=input_tensor) + print(output_shape) + + """ + + def __init__(self): + self.models = self._get_supported_models() + + def _get_supported_models(self) -> List[str]: + """Retrieve the list of supported models from timm.""" + return timm.list_models() + + def _create_model(self, model_info: TimmModelInfo) -> torch.nn.Module: + """ + Create a model instance from timm with specified parameters. + + Args: + model_info: An instance of TimmModelInfo containing model specifications. + + Returns: + An instance of a pytorch model. + """ + return timm.create_model( + model_info.model_name, + pretrained=model_info.pretrained, + in_chans=model_info.in_chans, + ) + + def __call__( + self, model_info: TimmModelInfo, input_tensor: torch.Tensor + ) -> torch.Size: + """ + Create and run a model specified by `model_info` on `input_tensor`. + + Args: + model_info: An instance of TimmModelInfo containing model specifications. + input_tensor: A torch tensor representing the input data. + + Returns: + The shape of the output from the model. + """ + model = self._create_model(model_info) + return model(input_tensor).shape diff --git a/swarms/models/zephyr.py b/swarms/models/zephyr.py index f4052d82..30d2bcd6 100644 --- a/swarms/models/zephyr.py +++ b/swarms/models/zephyr.py @@ -25,12 +25,22 @@ class Zephyr: def __init__( self, + model_name: str = "HuggingFaceH4/zephyr-7b-alpha", + tokenize: bool = False, + add_generation_prompt: bool = True, + system_prompt: str = "You are a friendly chatbot who always responds in the style of a pirate", max_new_tokens: int = 300, temperature: float = 0.5, top_k: float = 50, top_p: float = 0.95, + *args, + **kwargs, ): super().__init__() + self.model_name = model_name + self.tokenize = tokenize + self.add_generation_prompt = add_generation_prompt + self.system_prompt = system_prompt self.max_new_tokens = max_new_tokens self.temperature = temperature self.top_k = top_k @@ -38,14 +48,14 @@ class Zephyr: self.pipe = pipeline( "text-generation", - model="HuggingFaceH4/zephyr-7b-alpha", + model=self.model_name, torch_dtype=torch.bfloat16, device_map="auto", ) self.messages = [ { "role": "system", - "content": "You are a friendly chatbot who always responds in the style of a pirate", + "content": f"{self.systen_prompt}\n\nUser:", }, { "role": "user", @@ -53,10 +63,43 @@ class Zephyr: }, ] - def __call__(self, text: str): + def __call__(self, task: str): """Call the model""" prompt = self.pipe.tokenizer.apply_chat_template( - self.messages, tokenize=False, add_generation_prompt=True + self.messages, + tokenize=self.tokenize, + add_generation_prompt=self.add_generation_prompt, ) - outputs = self.pipe(prompt, max_new_token=self.max_new_tokens) + outputs = self.pipe(prompt) # max_new_token=self.max_new_tokens) print(outputs[0])["generated_text"] + + def chat(self, message: str): + """ + Adds a user message to the conversation and generates a chatbot response. + """ + # Add the user message to the conversation + self.messages.append({"role": "user", "content": message}) + + # Apply the chat template to format the messages + prompt = self.pipe.tokenizer.apply_chat_template( + self.messages, tokenize=False, add_generation_prompt=True + ) + + # Generate a response + outputs = self.pipe( + prompt, + max_new_tokens=self.max_new_tokens, + do_sample=True, + temperature=self.temperature, + top_k=self.top_k, + top_p=self.top_p, + ) + + # Extract the generated text + generated_text = outputs[0]["generated_text"] + + # Optionally, you could also add the chatbot's response to the messages list + # However, the below line should be adjusted to extract the chatbot's response only + # self.messages.append({"role": "bot", "content": generated_text}) + + return generated_text