|
import gradio as gr |
|
import boto3, json, os, wandb |
|
|
|
from dotenv import load_dotenv, find_dotenv |
|
_ = load_dotenv(find_dotenv()) |
|
|
|
aws_access_key_id = os.environ["AWS_ACCESS_KEY_ID"] |
|
aws_secret_access_key = os.environ["AWS_SECRET_ACCESS_KEY"] |
|
wandb_api_key = os.environ["WANDB_API_KEY"] |
|
|
|
config = { |
|
"max_tokens_to_sample": 300, |
|
"model": "anthropic.claude-v2", |
|
"temperature": 0, |
|
"top_k": 250, |
|
"top_p": 0.999, |
|
} |
|
|
|
bedrock_runtime = boto3.client( |
|
aws_access_key_id = aws_access_key_id, |
|
aws_secret_access_key = aws_secret_access_key, |
|
service_name = "bedrock-runtime", |
|
region_name = "us-west-2" |
|
) |
|
|
|
def wandb_log(prompt, completion): |
|
wandb.login(key = wandb_api_key) |
|
wandb.init(project = "bedrock-llm", config = config) |
|
wandb.log({"prompt": str(prompt), "completion": str(completion)}) |
|
wandb.finish() |
|
|
|
def invoke(prompt): |
|
if (prompt == ""): |
|
raise gr.Error("Prompt is required.") |
|
|
|
completion = "" |
|
|
|
try: |
|
body = json.dumps({"prompt": "Human: " + prompt + "Assistant: ", |
|
"max_tokens_to_sample": config["max_tokens_to_sample"], |
|
"temperature": config["temperature"], |
|
"top_k": config["top_k"], |
|
"top_p": config["top_p"], |
|
"stop_sequences": ["Human: "] |
|
}) |
|
modelId = config["model"] |
|
accept = "application/json" |
|
contentType = "application/json" |
|
|
|
response = bedrock_runtime.invoke_model(body = body, |
|
modelId = modelId, |
|
accept = accept, |
|
contentType = contentType) |
|
|
|
response_body = json.loads(response.get("body").read()) |
|
completion = response_body["completion"] |
|
except Exception as e: |
|
completion = e |
|
|
|
raise gr.Error(e) |
|
finally: |
|
wandb_log(prompt, completion) |
|
|
|
return completion |
|
|
|
description = """<a href='https://www.gradio.app/'>Gradio</a> UI using <a href='https://aws.amazon.com/bedrock/'>Amazon Bedrock</a> API |
|
with <a href='https://www.anthropic.com/'>Anthropic</a> Claude 2 foundation model. |
|
RAG evaluation via <a href='https://wandb.ai/bstraehle'>Weights & Biases</a>.""" |
|
|
|
gr.close_all() |
|
demo = gr.Interface(fn = invoke, |
|
inputs = [gr.Textbox(label = "Prompt", lines = 1)], |
|
outputs = [gr.Textbox(label = "Completion", lines = 1)], |
|
title = "Generative AI - LLM", |
|
description = description) |
|
demo.launch() |