Step-by-step guide to deploying a RAG Model in Snowpark Container Service (SPCS) with Matillion Data Productivity Cloud

USE ROLE ACCOUNTADMIN;

CREATE ROLE test_role;
GRANT ROLE test_role TO USER <<Your Username>>;

-- Grant necessary privileges --
GRANT BIND SERVICE ENDPOINT ON ACCOUNT TO ROLE test_role;

-- Create compute pool --
CREATE COMPUTE POOL GPU_NV_M_pool
  MIN_NODES = 1
  MAX_NODES = 1
  INSTANCE_FAMILY = GPU_NV_M;

-- Create warehouse --
CREATE OR REPLACE WAREHOUSE spcs_wh WITH
  WAREHOUSE_SIZE='X-SMALL'
  AUTO_SUSPEND = 180
  AUTO_RESUME = true
  INITIALLY_SUSPENDED = false;

-- Grant privileges --
GRANT USAGE ON COMPUTE POOL GPU_NV_M_pool TO ROLE test_role;
GRANT MONITOR ON COMPUTE POOL GPU_NV_M_pool TO ROLE test_role;
GRANT ALL ON WAREHOUSE spcs_wh TO ROLE test_role;

-- Create database --
CREATE DATABASE tutorial_db;
GRANT OWNERSHIP ON DATABASE tutorial_db TO ROLE test_role;

-- Create External Access Integration (to download Embedding model)
CREATE OR REPLACE NETWORK RULE HF_EGRESS_RULE
    MODE = EGRESS
    TYPE = HOST_PORT
    VALUE_LIST = ('huggingface.co','cdn-lfs-us-1.huggingface.co','cdn-lfs.huggingface.co');

CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION RAG_ACCESS_INTEGRATION
    ALLOWED_NETWORK_RULES = (HF_EGRESS_RULE)
    ENABLED = true;

USE ROLE test_role;
USE DATABASE tutorial_db;
USE WAREHOUSE spcs_wh;

CREATE SCHEMA IF NOT EXISTS data_schema;
CREATE IMAGE REPOSITORY IF NOT EXISTS llama2_rag_repository;
CREATE STAGE IF NOT EXISTS tutorial_stage
  DIRECTORY = ( ENABLE = true );

# Use the NVIDIA CUDA base image
FROM nvcr.io/nvidia/pytorch:23.06-py3

# Set the working directory in the container
WORKDIR /app

# Set environment variable for cuBLAS
ENV LLAMA_CUBLAS=1
# Set environment variable for llama -cpp-python installation
ENV CMAKE_ARGS="-DLLAMA_CUBLAS=on"

# Copy the necessary files into the container
COPY llama2_service_gpu.py /app/llama2_service_gpu.py
COPY llama-2-7b-chat.Q6_K.gguf /app/llama-2-7b-chat.Q6_K.gguf
COPY templates/ /app/templates/
COPY faiss_AiDoc/ /app/faiss_AiDoc/
COPY requirements.txt /app/requirements.txt

# Install Python dependencies
RUN pip3 install --upgrade pip && pip3 install -r requirements.txt

# Expose port 8080 outside of the container (if needed)
EXPOSE 8080

# Run llama2_service_gpu.py when the container launches
CMD ["python3", "./llama2_service_gpu.py"]

accelerate  
langchain  
langchain -community  
Flask  
llama -cpp-python  
transformers  
InstructorEmbedding  
sentence_transformers  
huggingface  
faiss-gpu

spec:  
  containers:                            
  - name: llama  
    image: /tutorial_db/data_schema/llama2_rag _repository  /llama2_gpu:llama2  
    env:                                 
      SYNC_DIR: /tmp/sync  
      SERVER_PORT: 8080     
    volumeMounts:                        
      - name: sync  
        mountPath: /tmp/sync  
      - name: shm  
        mountPath: /dev/shm  
    resources:                           
      requests:        
        nvidia.com/gpu: 3  
        cpu: 1  
      limits:  
        nvidia.com/gpu: 3  
  endpoints:                              
  - name: llamaendpoint  
    port: 8080  
    public: true  
  volumes:                                
  - name: sync  
    source: local  
  #  size: local  
  - name: shm  
    source: local

# Standard library imports
#!/usr/bin/env python
import os
import sys
import logging

# Third-party imports
from flask import Flask, jsonify, request, make_response, render_template
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.chains import LLMChain
from langchain_community.llms import LlamaCpp
from llama_cpp import Llama
from langchain.prompts import PromptTemplate
from langchain_community.embeddings import HuggingFaceInstructEmbeddings
from langchain_community.vectorstores import FAISS

# Set up environment variables
SERVICE_HOST = os.getenv('SERVER_HOST', '0.0.0.0')
SERVICE_PORT = os.getenv('SERVER_PORT', 8080)

# Set up logging
def get_logger(logger_name):
    logger = logging.getLogger(logger_name)
    logger.setLevel(logging.DEBUG)
    handler = logging.StreamHandler(sys.stdout)
    handler.setLevel(logging.DEBUG)
    handler.setFormatter(logging.Formatter('%(name)s [%(asctime)s] [%(levelname)s] %(message)s'))
    logger.addHandler(handler)
    return logger

# Configure the logger globally
logger = get_logger('llama-service')

# Initialize Flask app
app = Flask(__name__)

@app.get("/healthcheck")
def readiness_probe():
    # Added check to see if model and DB are loaded
    if model is not None and db is not None:
        return "Model and database are ready!"
    return "Waiting for model and database initialization..."

# Initialize and load the Llama model and FAISS database
model = None
db = None

def load_model():
    global model, db
    if model is None or db is None:
        logger.info("Initializing the Llama model and FAISS database...")  # Log initialization
        # Initialize the Llama model
        try:
            logger.info("test1")
            template = """Question: {question} \n\nAnswer:"""
            prompt = PromptTemplate(template=template, input_variables=["question"])

            # Callbacks support token-wise streaming
            callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])

            n_gpu_layers = 40
            n_batch = 512

            model_path = "llama-2-7b-chat.Q6_K.gguf"
            logger.info("test2")
            logger.info("Initializing LlamaCpp...")
            llm = LlamaCpp(
            model_path=model_path,
            n_gpu_layers=n_gpu_layers,
            n_ctx=4095,
            n_batch=n_batch,
            top_p=1,
            callback_manager=callback_manager,
            temperature = 0.4,
            max_tokens = 2095,
            n_parts=1,
            verbose=True
            )
            model = LLMChain(prompt=prompt, llm=llm)
            logger.info("Llama model initialized successfully.")  # Log successful initialization
        except Exception as e:
            logger.error(f"Error during Llama model initialization: {e}")
            logger.error(e, exc_info=True)  # Log the full exception traceback
            sys.exit(1)  # Exit the script

        # Initialize the FAISS database
        try:
            hf_embedding = HuggingFaceInstructEmbeddings()
            db_path = "faiss_AiDoc/"
            db = FAISS.load_local(db_path, embeddings=hf_embedding)
            logger.info("FAISS database loaded successfully.")
        except Exception as e:
            logger.error("Error during FAISS database initialization:")
            logger.error(e, exc_info=True)  # Log the full exception traceback
            sys.exit(1)  # Exit the script

    else:
        logger.info("Llama model and FAISS database already initialized.")  # Log if already initialized

    return model, db

@app.post("/llama")
def llama():
    message = request.json
    logger.debug(f'Received request: {message}')

    try:
        input_rows = message.get('data', [])

        if not input_rows:
            logger.info('Received empty message or no data key present')
            return make_response(jsonify({"data": []}), 200)

        if not all(isinstance(row, list) and len(row) >= 2 for row in input_rows):
            logger.error('Invalid input format')
            return make_response(jsonify({"error": "Invalid input format"}), 400)

        logger.info(f'Received {len(input_rows)} rows')
        output_rows = []
        for row in input_rows:
            try:
                response = get_llama_response(row[1])
                output_rows.append([row[0], response])
            except Exception as e:
                logger.error(f'Error processing row {row[0]}: {e}')
                output_rows.append([row[0], "Error processing input"])

        logger.info(f'Produced {len(output_rows)} rows')
        response = make_response(jsonify({"data": output_rows}), 200)
        response.headers['Content-type'] = 'application/json'
        logger.debug(f'Sending response: {response.get_json()}')
        return response
    except Exception as e:
        logger.error(f'Error in /llama endpoint: {e}')
        return make_response(jsonify({"error": "Internal Server Error"}), 500)

@app.route("/ui", methods=["GET", "POST"])
def ui():
    if request.method == "POST":
        input_text = request.form.get("input")
        return render_template("basic_ui.html", llama_input=input_text, llama_reponse=get_llama_response(input_text))
    return render_template("basic_ui.html")

def prompt_response(input):
    global model, db
    query = input
    search = db.similarity_search(query, k=5)

    template = '''Context: {context}

    Based on the Context, please answer the following question:
    Question: {question}

Provide an answer based on the context only, without using general knowledge. The answer should be derived directly from the context provided.
    Please correct any grammatical errors for improved readability.
    If the context does not contain relevant information to answer the question, state that the answer is not available in the given context.
    Please include the source title of the information as a reference of how you arrive at your answer. '''

    prompt = PromptTemplate(input_variables=["context", "question"], template=template)
    final_prompt = prompt.format(question=query, context=search)

    response = model.run(final_prompt)
    return response

def process_input(input):
    processed_answer = prompt_response(input)
    return processed_answer

def get_llama_response(input):
    reply = process_input(input)
    return f'{reply}'

if __name__ == '__main__':
    try:
        model, db = load_model()
        app.run(host=SERVICE_HOST, port=int(SERVICE_PORT), debug=False)
    except OSError as e:
        if e.errno == 98:
            logger.error("Address already in use. Please ensure the port is free or specify a different port.")
        else:
            raise

docker build --rm --platform linux/amd64 -t llama2_gpu:llama2

docker tag llama2_gpu:llama2 <<your Snowflake container repository URL>>

docker login myorg-myacct.registry.snowflakecomputing.com -u <<Username>>

docker push <<your Snowflake container repository URL>>  /llama2_gpu:llama2

CREATE SERVICE llama2_rag_service
    IN COMPUTE POOL GPU_NV_M_pool FROM @oae_llama2_stage_gpu
    SPEC='llama2.yaml'
    MIN_INSTANCES=1
    MAX_INSTANCES=1
    EXTERNAL_ACCESS_INTEGRATIONS = ( RAG_ACCESS_INTEGRATION );

CREATE SERVICE llama2_rag_service
  IN COMPUTE POOL GPU_NV_M_pool
  FROM SPECIFICATION $$
    spec:
     containers:
     - name: llama
       image: /tutorial_db/data_schema/  llama2_rag _repository  /llama2_gpu:llama2
       env:
         SYNC_DIR: /tmp/sync
         SERVER_PORT: 8080
      volumeMounts:
       - name: sync
         mountPath: /tmp/sync
       - name: shm
         mountPath: /dev/shm
      resources:
        requests:
          nvidia.com/gpu: 2
          cpu: 1
        limits:
          nvidia.com/gpu: 2
   endpoints:
   - name: llamaendpoint
     port: 8080

public: true
   volumes:
   - name: sync
     source: local
   - name: shm
    source: local
      $$
   MIN_INSTANCES=1
   MAX_INSTANCES=1
   EXTERNAL_ACCESS_INTEGRATIONS = ( RAG_ACCESS_INTEGRATION );

CREATE OR REPLACE FUNCTION llama2_rag (text varchar)
  RETURNS varchar
  SERVICE=  llama2_rag_service
  ENDPOINT=llamaendpoint
  AS '/llama';

SELECT llama2_rag('<<your question>>');