from google.cloud import aiplatform
import os

def deploy_hf_model(
    project_id: str,
    location: str,
    model_id: str,
    machine_type: str =  "a2-highgpu-4g,
):
    """
    Deploy a Hugging Face model using pre-built container
    """
    # Initialize Vertex AI
    aiplatform.init(project=project_id, location=location)
    
    env_vars = {
        "MODEL_ID": model_id,
        "MAX_INPUT_LENGTH": "512",
        "MAX_TOTAL_TOKENS": "1024",
        "MAX_BATCH_PREFILL_TOKENS": "2048",
        "NUM_SHARD": "1",
        # "HF_TOKEN": ""  # Add your Hugging Face token if needed
    }


    # Create model using pre-built container
    model = aiplatform.Model.upload(
        display_name=f"hf-{model_id.replace('/', '-')}",
        # Using the official container for Hugging Face models
        serving_container_image_uri="us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-2.ubuntu2204.py310",
        serving_container_environment_variables=env_vars
    )
  
    print(f"model uploaded: {model}")

    # Deploy model to endpoint
    endpoint = model.deploy(
        machine_type=machine_type,
        min_replica_count=1,
        max_replica_count=1,
        accelerator_type="NVIDIA_TESLA_A100",
        accelerator_count=1,
        sync=True
    )
    
    print(f"Model deployed to endpoint: {endpoint.resource_name}")
    return endpoint


def create_completion(
    endpoint,
    prompt: str,
    max_tokens: int = 100,
    temperature: float = 0.7
):
    """
    Generate text using deployed model
    """
    response = endpoint.predict({
        "text": prompt,
        "parameters": {
            "max_new_tokens": max_tokens,
            "temperature": temperature,
            "top_p": 0.95,
            "top_k": 40,
        }
    })
    return response

if __name__ == "__main__

bashCopy# List endpoints
gcloud ai endpoints list

# Get endpoint details
gcloud ai endpoints describe ENDPOINT_ID

# Get endpoint predictions
gcloud ai endpoints predict ENDPOINT_ID \
    --json-request='{"instances": [{"text": "Tell me a joke"}]