Models

import requests
import json

# Submit a managed fine-tuning job via the Databricks REST API
DATABRICKS_HOST = "https://your-workspace.azuredatabricks.net"
TOKEN = "dapiXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"

payload = {
    "model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
    "train_data_path": "dbfs:/user/data/instruction_pairs_train",
    "eval_data_path": "dbfs:/user/data/instruction_pairs_eval",
    "training_duration": "3ep",           # 3 epochs
    "learning_rate": 5e-6,
    "register_to": "ml_catalog.llm_models.llama_support_ft",
    "data_prep_cluster_id": "0101-123456-abcdefgh",
}

# Production code should wrap this in try/except for network errors.
response = requests.post(
    f"{DATABRICKS_HOST}/api/2.0/fine-tuning/runs/create",
    headers={"Authorization": f"Bearer {TOKEN}"},
    json=payload,
)
run_id = response.json()["run_id"]
print(f"Fine-tuning run started: {run_id}")

from composer import Trainer
from composer.models import HuggingFaceModel
from composer.optim import DecoupledAdamW
from streaming import StreamingDataset
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "meta-llama/Meta-Llama-3.1-8B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
base_model = AutoModelForCausalLM.from_pretrained(
    model_name, torch_dtype=torch.bfloat16
)

# Wrap in Composer's HuggingFace adapter
composer_model = HuggingFaceModel(
    model=base_model,
    tokenizer=tokenizer,
    use_logits=True,
)

# Streaming dataset from DBFS (MDS format written by Spark)
train_dataset = StreamingDataset(
    remote="s3://my-bucket/training-data/mds/",
    local="/tmp/streaming_cache/",
    shuffle=True,
    batch_size=8,
)

trainer = Trainer(
    model=composer_model,
    train_dataloader=torch.utils.data.DataLoader(train_dataset, batch_size=8),
    max_duration="3ep",
    optimizers=DecoupledAdamW(composer_model.parameters(), lr=5e-6, weight_decay=0.0),
    device="gpu",
    precision="bf16_mixed",
    save_folder="/dbfs/user/checkpoints/llama_sft/",
    loggers=[],  # MLflow logger added via mlflow.start_run() context
)
trainer.fit()

from openai import OpenAI

# The openai client works with Databricks Foundation Model APIs
# by overriding the base_url and using a Databricks PAT as the key
client = OpenAI(
    api_key="dapiXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX",
    base_url="https://your-workspace.azuredatabricks.net/serving-endpoints",
)

response = client.chat.completions.create(
    model="databricks-dbrx-instruct",   # or databricks-meta-llama-3-1-70b-instruct
    messages=[
        {"role": "system", "content": "You are a helpful data engineering assistant."},
        {"role": "user", "content": "Explain the difference between Delta Live Tables and standard Delta tables."},
    ],
    max_tokens=512,
    temperature=0.3,
)
print(response.choices[0].message.content)

-- Classify customer support tickets by category using ai_classify()
SELECT
    ticket_id,
    ticket_text,
    ai_classify(
        ticket_text,
        ARRAY('billing', 'technical', 'account_access', 'feature_request', 'other')
    ) AS category,
    ai_extract(
        ticket_text,
        ARRAY('product_name', 'error_code', 'urgency_level')
    ) AS extracted_fields,
    ai_summarize(ticket_text, 50) AS short_summary
FROM raw_support_tickets
WHERE created_date >= CURRENT_DATE - INTERVAL 1 DAY;

-- Use ai_query() for custom prompts inside a Delta Live Tables pipeline
CREATE OR REFRESH STREAMING TABLE product_reviews_enriched AS
SELECT
    review_id,
    product_id,
    review_text,
    rating,
    -- Call the Foundation Model API for sentiment and entity extraction
    ai_query(
        'databricks-meta-llama-3-1-70b-instruct',
        CONCAT(
            'Analyze this product review. Return a JSON object with keys: ',
            '"sentiment" (positive/negative/neutral), ',
            '"mentioned_features" (array of product features mentioned), ',
            '"improvement_suggestions" (array of suggestions, empty if none). ',
            'Review: ', review_text
        )
    ) AS llm_analysis,
    -- Also classify sentiment directly
    ai_classify(review_text, ARRAY('positive', 'negative', 'neutral')) AS sentiment_label
FROM STREAM(LIVE.raw_product_reviews);

import mlflow
import mlflow.openai
from mlflow.models import infer_signature
import pandas as pd

mlflow.set_experiment("/Users/you@company.com/llm-finetuning-v3")

with mlflow.start_run(run_name="llama-support-ft-eval"):
    # Log hyperparameters
    mlflow.log_params({
        "base_model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
        "learning_rate": 5e-6,
        "training_epochs": 3,
        "train_samples": 12400,
    })

    # Log prompt template used for fine-tuning as an artifact
    mlflow.log_text(
        "You are a helpful support agent. Answer the customer's question concisely.\n\nQuestion: {question}\n\nAnswer:",
        "prompt_template.txt",
    )

    # Evaluate the fine-tuned model with mlflow.evaluate()
    eval_data = pd.DataFrame({
        "inputs": ["How do I reset my password?", "Where is my order?"],
        "ground_truth": ["Go to Settings > Security > Reset Password.", "Check your order status at orders.example.com."],
    })

    results = mlflow.evaluate(
        model="endpoints:/ml_catalog.llm_models.llama_support_ft/1",
        data=eval_data,
        targets="ground_truth",
        model_type="question-answering",
        evaluators="default",
        extra_metrics=[
            mlflow.metrics.faithfulness(),
            mlflow.metrics.answer_relevance(),
            mlflow.metrics.toxicity(),
        ],
    )

    # Log aggregate metrics from the evaluation
    mlflow.log_metrics({
        "faithfulness_mean": results.metrics["faithfulness/v1/mean"],
        "relevance_mean": results.metrics["answer_relevance/v1/mean"],
        "toxicity_ratio": results.metrics["toxicity/v1/ratio"],
    })

    print(results.tables["eval_results_table"])

from mlflow.tracking import MlflowClient

client = MlflowClient()
model_name = "ml_catalog.llm_models.llama_support_ft"

# Transition version 3 to Champion (replaces the previous champion)
client.set_registered_model_alias(
    name=model_name,
    alias="Champion",
    version="3",
)

# Create or update a Model Serving endpoint to use the new champion
import requests, json

DATABRICKS_HOST = "https://your-workspace.azuredatabricks.net"
TOKEN = "dapiXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"

endpoint_config = {
    "name": "support-llm-endpoint",
    "config": {
        "served_models": [{
            "name": "llama-support-ft-champion",
            "model_name": model_name,
            "model_version": "3",
            "workload_size": "Small",
            "scale_to_zero_enabled": True,
        }]
    }
}

requests.put(
    f"{DATABRICKS_HOST}/api/2.0/serving-endpoints/support-llm-endpoint/config",
    headers={"Authorization": f"Bearer {TOKEN}"},
    json=endpoint_config["config"],
)

from databricks.vector_search.client import VectorSearchClient

vsc = VectorSearchClient()

# Create a vector search endpoint (compute that serves queries)
vsc.create_endpoint(
    name="llm-rag-endpoint",
    endpoint_type="STANDARD",
)

# Create a Delta Sync index: Databricks computes embeddings automatically
# using the specified embedding model endpoint, and re-syncs when the
# source Delta table changes
index = vsc.create_delta_sync_index(
    endpoint_name="llm-rag-endpoint",
    source_table_name="ml_catalog.rag_data.product_docs",
    index_name="ml_catalog.rag_data.product_docs_index",
    pipeline_type="TRIGGERED",           # or CONTINUOUS for near-real-time sync
    primary_key="doc_id",
    embedding_source_column="content",   # column containing text to embed
    embedding_model_endpoint_name="databricks-bge-large-en",
)

# Query the vector search index at inference time
results = index.similarity_search(
    query_text="How do I configure SSO with Azure AD?",
    columns=["doc_id", "title", "content", "last_updated"],
    num_results=5,
    filters={"doc_type": "integration_guide"},  # optional metadata filter
)

for hit in results["result"]["data_array"]:
    doc_id, title, content, last_updated, score = hit
    print(f"[{score:.3f}] {title} ({doc_id})")

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, udf, monotonically_increasing_id, explode
from pyspark.sql.types import ArrayType, StructType, StructField, StringType, IntegerType

spark = SparkSession.builder.getOrCreate()

# UDF: split document text into overlapping chunks
def chunk_text(text: str, chunk_size: int = 512, overlap: int = 64):
    """Split text into chunks of ~chunk_size characters with overlap."""
    if not text:
        return []
    chunks = []
    start = 0
    while start < len(text):
        end = min(start + chunk_size, len(text))
        chunks.append({"chunk_text": text[start:end], "chunk_index": len(chunks)})
        if end == len(text):
            break
        start += chunk_size - overlap
    return chunks

chunk_schema = ArrayType(StructType([
    StructField("chunk_text", StringType()),
    StructField("chunk_index", IntegerType()),
]))
chunk_udf = udf(chunk_text, chunk_schema)

# Explode raw documents into chunk rows
raw_docs = spark.table("ml_catalog.rag_data.raw_documents")

chunked = (
    raw_docs
    .withColumn("chunks", chunk_udf(col("content")))
    .select("doc_id", "title", "doc_type", "last_updated", explode("chunks").alias("chunk"))
    .select(
        "doc_id", "title", "doc_type", "last_updated",
        col("chunk.chunk_text").alias("content"),
        col("chunk.chunk_index").alias("chunk_index"),
    )
    .withColumn("chunk_id", monotonically_increasing_id().cast("string"))
)

chunked.write.format("delta").mode("append").saveAsTable("ml_catalog.rag_data.product_docs")

import mlflow
import mlflow.pyfunc
from databricks.vector_search.client import VectorSearchClient
from openai import OpenAI

class DatabricksRAGChain(mlflow.pyfunc.PythonModel):
    """MLflow pyfunc wrapper for a Databricks-native RAG chain."""

    def load_context(self, context):
        self.vsc = VectorSearchClient()
        self.index = self.vsc.get_index(
            endpoint_name="llm-rag-endpoint",
            index_name="ml_catalog.rag_data.product_docs_index",
        )
        self.llm = OpenAI(
            api_key=context.model_config["databricks_token"],
            base_url=context.model_config["databricks_host"] + "/serving-endpoints",
        )
        self.model_name = context.model_config.get(
            "generation_model", "databricks-dbrx-instruct"
        )

    def predict(self, context, model_input):
        questions = model_input["question"].tolist()
        answers = []

        for question in questions:
            # Retrieve top-5 relevant chunks
            results = self.index.similarity_search(
                query_text=question,
                columns=["title", "content"],
                num_results=5,
            )
            context_text = "\n\n".join(
                f"[{row[0]}]\n{row[1]}"
                for row in results["result"]["data_array"]
            )

            # Generate answer with retrieved context
            response = self.llm.chat.completions.create(
                model=self.model_name,
                messages=[
                    {
                        "role": "system",
                        "content": (
                            "Answer the question using only the provided context. "
                            "If the context does not contain the answer, say so."
                        ),
                    },
                    {
                        "role": "user",
                        "content": f"Context:\n{context_text}\n\nQuestion: {question}",
                    },
                ],
                max_tokens=512,
                temperature=0.1,
            )
            answers.append(response.choices[0].message.content)

        import pandas as pd
        return pd.DataFrame({"answer": answers})

# Log the chain to MLflow and register it
with mlflow.start_run(run_name="rag-chain-v1"):
    model_config = {
        "databricks_host": "https://your-workspace.azuredatabricks.net",
        "databricks_token": "{{secrets/rag-scope/databricks-token}}",
        "generation_model": "databricks-dbrx-instruct",
    }
    mlflow.pyfunc.log_model(
        artifact_path="rag_chain",
        python_model=DatabricksRAGChain(),
        model_config=model_config,
        registered_model_name="ml_catalog.rag_apps.product_docs_rag",
    )