TinyLLaVA로 LLM 앱 구축

!pip install wandb -qqq
!pip install git+https://github.com/huggingface/transformers -qqq
!pip install --upgrade gradio -qqq
!pip install chromadb -qqq
!pip install openai -qqq
!pip install tiktoken -qqq
!pip install langchain -qqq
!pip install accelerate -qqq
!pip install bitsandbytes -qqq

import os
import requests
import numpy as np
import torch
import datetime

# For loading in the tiny-LLaVA-v1-hf model in a transformers pipeline.
import transformers
from transformers import pipeline
from transformers import BitsAndBytesConfig

# For converting input images to PIL images.
from PIL import Image

# For creating the gradio app.
import gradio as gr

# For creating a simple prompt (open to extension) to our model.
from langchain.prompts import PromptTemplate

# Our vector database of choice: Chroma!
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings.openai import OpenAIEmbeddings

import chromadb
from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction
from chromadb.utils.data_loaders import ImageLoader

# For loading in our OpenAI API key.
from google.colab import userdata

# For logging.
import wandb
from wandb.sdk.data_types.trace_tree import Trace
wandb.login()

# Required for us to load in our pipeline for TinyLLaVA.
assert transformers.__version__ >= "4.35.3"

# Use OpenAI's embeddings for our Chroma collection.
embeddings = OpenAIEmbeddings(
    model="text-embedding-ada-002",
    openai_api_key=userdata.get("OPENAI_API_KEY"),
)
collection = Chroma("conversation_memory", embeddings)

# Ref: https://huggingface.co/bczhou/tiny-llava-v1-hf
model_id = "bczhou/tiny-llava-v1-hf"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

pipe = pipeline(
    "image-to-text",
    model=model_id,
    device_map="auto",
    use_fast=True,
    model_kwargs={"quantization_config": bnb_config}
)

try:
  assert user_avatar_image_path
except:
  img_data = requests.get("https://imgur.com/QehpHeV.png").content
  with open('user_avatar.png', 'wb') as handler:
      handler.write(img_data)
  user_avatar_image_path = "user_avatar.png"

try:
  assert chatbot_avatar_image_path
except:
  img_data = requests.get("https://imgur.com/ki4hPhZ.png").content
  with open('chatbot_avatar.png', 'wb') as handler:
      handler.write(img_data)
  chatbot_avatar_image_path = "chatbot_avatar.png"

# Let's get a sample image to use. You can download it and pass it into the app!
# The prompt is: What's unusual about this image?
img_data = requests.get("https://imgur.com/Ca6gjuf.png").content
with open('sample_image.png', 'wb') as handler:
    handler.write(img_data)

max_new_tokens = 200

# Path for storing images.
IMG_ROOT_PATH = "data/"
os.makedirs(IMG_ROOT_PATH, exist_ok=True)

def generate_output(message: str, history: list, img: np.ndarray) -> str:
    """Generates an output given a message and image."""

    status = "success"

    # Get detailed description of the image for Chroma.
    query = "Please provide a detailed description of the image."
    prompt = PromptTemplate.from_template(
        "USER: <image>\n" +
        "{query}" +
        "\n" +
        "ASSISTANT: "
    )

    start_time_ms = datetime.datetime.now().timestamp() * 1000
    try:
        outputs = pipe(Image.fromarray(img), prompt=prompt.format(query=query), generate_kwargs={"max_new_tokens": max_new_tokens})
        img_desc = outputs[0]["generated_text"].split("ASSISTANT:")[-1]
        status_message = (None,)
    except Exception as e:
        status = "error"
        status_message = str(e)
        img_desc = ""
    end_time_ms = round(datetime.datetime.now().timestamp() * 1000)

    # Create a span in wandb.
    root_span = Trace(
        name="img_desc_span",
        kind="llm",  # kind can be "llm", "chain", "agent" or "tool"
        status_code=status,
        status_message=status_message,
        metadata={
            "max_new_tokens": max_new_tokens,
            "model_name": model_id,
        },
        start_time_ms=start_time_ms,
        end_time_ms=end_time_ms,
        inputs={"system_prompt": prompt.format(query=query), "query": query},
        outputs={"response": img_desc},
    )

    # Log the span to wandb.
    root_span.log(name="img_desc_trace") 

    # Visual Question-Answering!
    prompt = PromptTemplate.from_template(
        "Context: {context}\n\n"
        "USER: <image>\n" +
        "{message}" +
        "\n" +
        "ASSISTANT: "
    )
    context = collection.similarity_search(query=message, k=2)
    context = "\n".join([doc.page_content for doc in context])

    # Forward pass through the model with given prompt template.
    start_time_ms = datetime.datetime.now().timestamp() * 1000
    try:
        outputs = pipe(
            Image.fromarray(img),
            prompt=prompt.format(
                context=context,
                message=message
            ),
            generate_kwargs={"max_new_tokens": max_new_tokens}
        )
        response = outputs[0]["generated_text"].split("ASSISTANT:")[-1]
        status_message = (None,)
    except Exception as e:
      status = "error"
      status_message = str(e)
      response = ""
    end_time_ms = round(datetime.datetime.now().timestamp() * 1000)

    # Create a span in wandb.
    root_span = Trace(
        name="response_span",
        kind="llm",  # kind can be "llm", "chain", "agent" or "tool"
        status_code=status,
        status_message=status_message,
        metadata={
            "max_new_tokens": max_new_tokens,
            "model_name": model_id,
        },
        start_time_ms=start_time_ms,
        end_time_ms=end_time_ms,
        inputs={
            "system_prompt": prompt.format(
                context=context,
                message=message
            ),
            "query": message
        },
        outputs={"response": response},
    )

    # Log the span to wandb.
    root_span.log(name="response_trace")

    # Add (img_desc, message, response) 3-tuple to Chroma collection.
    text = f"Image Description: {img_desc}\nUSER: {message}\nASSISTANT: {response}\n"
    collection.add_texts(texts=[text])

    # Return model output.
    return img_desc + "\n\n" + response

wandb.init(project="building_llm_app")

# Define the ChatInterface, customize, and launch!
gr.ChatInterface(
    generate_output,
    chatbot=gr.Chatbot(
        label="Chat with me!",
        show_label=True,
        container=False,
        scale=5,
        height=300,
        show_share_button=True,
        show_copy_button=True,
        avatar_images=(user_avatar_image_path, chatbot_avatar_image_path),
        likeable=False,
        layout="bubble",
        bubble_full_width=False
      ),
    textbox=gr.Textbox(
        lines=1,
        max_lines=5,
        placeholder="Message ...",
        container=False,
        scale=7,
        info="Input your textual response in the text field and your image below!"
    ),
    additional_inputs="image",
    additional_inputs_accordion=gr.Accordion(
        open=True,
    ),
    title="Language-Image Question Answering with bczhou/TinyLLaVA-v1-hf!",
    description="""
    This simple gradio app internally uses a Large Language-Vision Model (LLVM) and the Chroma vector database for memory.
    Note: this minimal app requires both an image and a text-based query before the chatbot system can respond.
    """,
    theme="soft",
    submit_btn="Submit ▶",
    retry_btn=None,
    undo_btn="Delete Previous",
    clear_btn="Clear",
).launch(debug=True, share=True)