Vertex AI

# Vertex AI를 활용한 Gemini API 및 MLOps 예제

from google.cloud import aiplatform
from vertexai.generative_models import GenerativeModel, Part
from vertexai.language_models import TextEmbeddingModel
import vertexai

# 프로젝트 초기화
PROJECT_ID = "my-gcp-project"
LOCATION = "us-central1"
vertexai.init(project=PROJECT_ID, location=LOCATION)

# ==========================================
# 1. Gemini API 사용 (멀티모달)
# ==========================================

def chat_with_gemini(prompt: str, image_path: str = None) -> str:
    """Gemini 모델로 대화형 응답 생성"""
    model = GenerativeModel("gemini-1.5-pro")

    if image_path:
        # 이미지와 함께 멀티모달 입력
        image = Part.from_uri(image_path, mime_type="image/jpeg")
        response = model.generate_content([prompt, image])
    else:
        response = model.generate_content(prompt)

    return response.text

# Function Calling 예제
def gemini_with_function_calling():
    """Gemini의 함수 호출 기능"""
    from vertexai.generative_models import FunctionDeclaration, Tool

    # 함수 정의
    get_weather_func = FunctionDeclaration(
        name="get_weather",
        description="지정된 도시의 현재 날씨 정보를 조회합니다",
        parameters={
            "type": "object",
            "properties": {
                "city": {"type": "string", "description": "도시 이름"},
                "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}
            },
            "required": ["city"]
        }
    )

    weather_tool = Tool(function_declarations=[get_weather_func])
    model = GenerativeModel("gemini-1.5-pro", tools=[weather_tool])

    response = model.generate_content("서울 날씨 알려줘")

    # 함수 호출 결과 처리
    if response.candidates[0].content.parts[0].function_call:
        func_call = response.candidates[0].content.parts[0].function_call
        print(f"Function: {func_call.name}, Args: {func_call.args}")

    return response

# ==========================================
# 2. 텍스트 임베딩 생성
# ==========================================

def create_embeddings(texts: list) -> list:
    """텍스트 임베딩 벡터 생성 (RAG 구축용)"""
    model = TextEmbeddingModel.from_pretrained("text-embedding-004")

    embeddings = model.get_embeddings(texts)
    vectors = [emb.values for emb in embeddings]

    return vectors

# ==========================================
# 3. Custom Training Job
# ==========================================

def run_custom_training():
    """커스텀 학습 작업 실행"""
    aiplatform.init(project=PROJECT_ID, location=LOCATION)

    # 커스텀 학습 작업 생성
    job = aiplatform.CustomTrainingJob(
        display_name="my-custom-training",
        script_path="trainer/train.py",  # 학습 스크립트
        container_uri="us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-13:latest",
        requirements=["transformers", "datasets"],
        model_serving_container_image_uri="us-docker.pkg.dev/vertex-ai/prediction/pytorch-gpu.1-13:latest"
    )

    # 학습 실행
    model = job.run(
        dataset=None,  # 또는 Vertex AI Dataset 참조
        model_display_name="trained-model",
        machine_type="n1-standard-8",
        accelerator_type="NVIDIA_TESLA_T4",
        accelerator_count=1,
        args=["--epochs", "10", "--batch-size", "32"]
    )

    return model

# ==========================================
# 4. 모델 배포 및 예측
# ==========================================

def deploy_model(model_resource_name: str):
    """학습된 모델을 엔드포인트에 배포"""
    model = aiplatform.Model(model_resource_name)

    # 엔드포인트 생성 및 배포
    endpoint = model.deploy(
        deployed_model_display_name="production-model",
        machine_type="n1-standard-4",
        min_replica_count=1,
        max_replica_count=3,
        accelerator_type="NVIDIA_TESLA_T4",
        accelerator_count=1,
        traffic_split={"0": 100}  # 100% 트래픽
    )

    return endpoint

def predict(endpoint, instances: list):
    """배포된 모델로 예측"""
    predictions = endpoint.predict(instances=instances)
    return predictions.predictions

# ==========================================
# 5. Vertex AI Pipelines (MLOps)
# ==========================================

from kfp import dsl
from kfp.v2 import compiler
from google.cloud.aiplatform import pipeline_jobs

@dsl.component(base_image="python:3.10")
def preprocess_data(input_path: str, output_path: str):
    """데이터 전처리 컴포넌트"""
    import pandas as pd
    df = pd.read_csv(input_path)
    # 전처리 로직
    df.to_csv(output_path, index=False)

@dsl.component(base_image="us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-13:latest")
def train_model(data_path: str, model_path: str, epochs: int):
    """모델 학습 컴포넌트"""
    # 학습 로직
    pass

@dsl.pipeline(name="ml-training-pipeline")
def ml_pipeline(input_data: str):
    """ML 파이프라인 정의"""
    preprocess_task = preprocess_data(
        input_path=input_data,
        output_path="gs://my-bucket/processed_data.csv"
    )

    train_task = train_model(
        data_path=preprocess_task.outputs["output_path"],
        model_path="gs://my-bucket/model/",
        epochs=10
    )

# 파이프라인 컴파일 및 실행
compiler.Compiler().compile(
    pipeline_func=ml_pipeline,
    package_path="pipeline.json"
)

job = pipeline_jobs.PipelineJob(
    display_name="training-run-001",
    template_path="pipeline.json",
    parameter_values={"input_data": "gs://my-bucket/raw_data.csv"}
)
job.submit()