Skip to content

Commit

Permalink
fuyu + zephyr fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
kyegomez committed Nov 8, 2023
1 parent 6b4c2d4 commit 9aa4084
Show file tree
Hide file tree
Showing 7 changed files with 136 additions and 41 deletions.
43 changes: 13 additions & 30 deletions demos/accountant_team/accountant_team.py
Original file line number Diff line number Diff line change
@@ -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}
"""


Expand All @@ -47,21 +40,11 @@ def summary_agent_prompt(analyzed_doc: str):
# 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}")
2 changes: 1 addition & 1 deletion playground/models/anthropic_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,6 @@
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))
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 <kye@apac.ai>"]
Expand Down
3 changes: 2 additions & 1 deletion swarms/models/fastvit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 4 additions & 3 deletions swarms/models/fuyu.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
"""Fuyu model by Kye"""
from transformers import (
FuyuProcessor,
FuyuForCausalLM,
AutoTokenizer,
FuyuProcessor,
FuyuImageProcessor,
)
from PIL import Image
Expand Down Expand Up @@ -50,9 +50,9 @@ def __init__(
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
)
Expand All @@ -62,3 +62,4 @@ def __call__(self, text: str, img_path: str):

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))
67 changes: 67 additions & 0 deletions swarms/models/timm.py
Original file line number Diff line number Diff line change
@@ -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
53 changes: 48 additions & 5 deletions swarms/models/zephyr.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,38 +25,81 @@ 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
self.top_p = top_p

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",
"content": "How many helicopters can a human eat in one sitting?",
},
]

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

0 comments on commit 9aa4084

Please sign in to comment.