Getting Started with Amazon Bedrock and W&B Weave

!pip install weave

import weave
weave.init('bedrock-weave')

import json
import boto3

@weave.op() # <- just add this 😎
def generate_text(
    model_id: str, 
    prompt: str, 
    max_tokens: int=400,
    temperature: float=0.7,
) -> dict:
    # Check your region model access page, in my case on us-east-1 region:
    # https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/modelaccess
    bedrock = boto3.client(service_name='bedrock-runtime')

    body = json.dumps({
        "prompt": prompt,
        "max_tokens": max_tokens,
        "temperature": temperature,
    })
    
    response = bedrock.invoke_model(body=body, modelId=model_id)

    response_body = json.loads(response.get('body').read())
    return response_body.get('outputs')

model_id = 'mistral.mistral-7b-instruct-v0:2'

# this is the required prompt format for mistral instruct
prompt = """<s>[INST] Create a Vegan Carbonara Recipe[/INST] """

outputs = generate_text(model_id, prompt)

for index, output in enumerate(outputs):

    print(f"Output {index + 1}\n----------")
    print(f"Text:\n{output['text']}\n")
    print(f"Stop reason: {output['stop_reason']}\n")

# Output 1
# ----------
# Text:
#  Vegan Carbonara Recipe:

# Ingredients:
# - 12 oz (340g) spaghetti or other long pasta
# - 1/2 cup (120ml) unsweetened almond milk or cashew cream
# - 1/2 cup (100g) nutritional yeast
# - 1/3 cup (80g) cooked and mashed sweet potato or pumpkin
# - 1 tbsp olive oil
# - 1 cup (200g) steamed or roasted vegetables such as zucchini, broccoli, or asparagus
# - 1/2 cup (100g) chopped mushrooms
# - 1/2 cup (100g) chopped onion
# - 1 clove garlic, minced
# - 1/2 cup (100g) chopped walnuts or cashews, toasted
# - 1/2 cup (100g) vegetable bacon or pancetta, cooked and crumbled (optional)
# - Salt and freshly ground black pepper, to taste
# - Red pepper flakes, for serving (optional)
# - Fresh parsley or basil, for serving (optional)

# Instructions:

# 1. Cook the pasta according to package instructions in salted water until al dente. Drain and set aside, reserving 1 cup of the pasta water.

# 2. In a blender or food processor, combine the almond milk or cashew cream, nutritional yeast, mashed sweet potato or pumpkin, and olive oil. Blend until smooth and creamy. Taste and adjust seasoning with salt and pepper as needed.

# 3. In a large skillet, sauté the chopped onion, garlic, and

# Stop reason: length

@weave.op()
def format_prompt(prompt: str) -> str:
    return f"""<s>[INST] {prompt}[/INST] """ 

@weave.op()
def run(prompt: str) -> None:
    prompt = format_prompt(prompt)

    outputs = generate_text(model_id, prompt, max_tokens=1000)
    
    for index, output in enumerate(outputs):

        print(f"Output {index + 1}\n----------")
        print(f"Text:\n{output['text']}\n")
        print(f"Stop reason: {output['stop_reason']}\n")
        
run("What clothes do I need to pack in winter for a trip to Chamonix?")

model_id = 'mistral.mixtral-8x7b-instruct-v0:1'

prompt = """<s>[INST] Create a Vegan Carbonara Recipe[/INST] """

outputs = generate_text(model_id, prompt, max_tokens=1000)

for index, output in enumerate(outputs):

    print(f"Output {index + 1}\n----------")
    print(f"Text:\n{output['text']}\n")
    print(f"Stop reason: {output['stop_reason']}\n")

@weave.op() # <- just add this 😎
def anthropic_call(
    model_id: str, 
    messages: str, 
    max_tokens: int=400,
    system_prompt: str=None,
) -> dict:
    bedrock = boto3.client(service_name='bedrock-runtime')


    body = json.dumps({
        "system": system_prompt,
        "messages": messages,
        "max_tokens": max_tokens,
        "anthropic_version": "bedrock-2023-05-31"
    })
    
    response = bedrock.invoke_model(
        body=body,
        modelId=model_id
    )

    response_body = json.loads(response.get('body').read())

    outputs = response_body.get('content')
    
    return outputs

from pathlib import Path
from dataclasses import dataclass

system_prompt = """
# Instructions

You are a documentation translation assistant from English to {output_language}. We are translating valid docusaurus flavored markdown. Some rules to remember:

- Do not add extra blank lines.
- The results must be valid docusaurus markdown
- It is important to maintain the accuracy of the contents but we don't want the output to read like it's been translated. So instead of translating word by word, prioritize naturalness and ease of communication.
- In code blocks, just translate the comments and leave the code as is.

## Formatting Rules

Do not translate target markdown links. Never translate the part of the link inside (). For instance here [https://wandb.ai/site](https://wandb.ai/site) do not translate anything, but on this, you should translate the [] part:
[track metrics](./guides/track), [create logs](./guides/artifacts).
"""

human_prompt = """
Here is a chunk of documentation in docusaurus Markdown format to translate. Return the translation only, without adding anything else. 
<Markdown start>
{md_chunk}
<End of Markdown>
"""

@dataclass
class PromptTemplate:
    system_prompt: str
    human_prompt: str
    language: str
    
    @weave.op() # <- just add this 😎
    def format_claude(self, md_chunk):
        "A formatting function for Mistral Instruct models"
        system_prompt = self.system_prompt.format(output_language=self.language)
        human_prompt = self.human_prompt.format(md_chunk=md_chunk)
        messages = [{"role":"user", "content":human_prompt}]
        return system_prompt, messages

from weave import Model

class ClaudeDocTranslator(Model):
    model_id: str='anthropic.claude-3-sonnet-20240229-v1:0'
    max_tokens: int=2048
    prompt_template: PromptTemplate
    
    @weave.op()
    def format_doc(self, path: Path) -> str:
        "Read and format the document"
        doc = Path(path).read_text()
        system_prompt, messages = prompt_template.format_claude(doc)
        return system_prompt, messages
    
    @weave.op()
    def translate(self, path: Path) -> dict:
        system_prompt, messages = self.format_doc(path)
        output = anthropic_call(
            self.model_id, 
            messages=messages,
            system_prompt=system_prompt,
            max_tokens=self.max_tokens)
        return output[0]['text']

model = ClaudeDocTranslator(prompt_template=prompt_template)
out = model.translate(Path("./docs/quickstart.md"))
print(out[0:400])