A guide to using the Azure AI model inference API


import weave
from azure.ai.inference import ChatCompletionsClient
from azure.ai.inference.models import SystemMessage, UserMessage
from azure.core.credentials import AzureKeyCredential

# Initialize Weave for logging
weave.init('azure-api')

# Define the endpoint URL and API key (hardcoded for this example)
ENDPOINT_URL = "your endpoint url"
PRIMARY_KEY = "your primary key"


PROMPT = "How many languages are in the world?"



# Initialize the ChatCompletionsClient
client = ChatCompletionsClient(
    endpoint=ENDPOINT_URL,
    credential=AzureKeyCredential(PRIMARY_KEY)
)

@weave.op
def run_inference(prompt, client):
    """
    Function to perform inference using the provided client and prompt.
    """
    response = client.complete(
        messages=[
            SystemMessage(content="You are a helpful assistant."),
            UserMessage(content=prompt),
        ]
    )
    try:
        content = response.choices[0].message.content
        print("Response:", content)
        return content
    except Exception as e:
        print(f"Failed to get response: {e}")
        return None


# Perform inference using the hardcoded prompt
run_inference(PROMPT, client)


python inference.py

from azure.ai.inference import ChatCompletionsClient
from azure.ai.inference.models import SystemMessage, UserMessage
from azure.core.credentials import AzureKeyCredential
import weave

# Initialize Weave for logging
weave.init('azure-api')

# Define global variables for the two models
MODEL_1_ENDPOINT_URL = "your first model endpoint url"
MODEL_1_PRIMARY_KEY = "your first model key"

MODEL_2_ENDPOINT_URL = "your second model endpoint url"
MODEL_2_PRIMARY_KEY = "your second model key"

PROMPT = "How many languages are in the world?"

# Initialize the ChatCompletionsClient for each model
client1 = ChatCompletionsClient(
    endpoint=MODEL_1_ENDPOINT_URL,
    credential=AzureKeyCredential(MODEL_1_PRIMARY_KEY)
)

client2 = ChatCompletionsClient(
    endpoint=MODEL_2_ENDPOINT_URL,
    credential=AzureKeyCredential(MODEL_2_PRIMARY_KEY)
)


@weave.op
def get_model_prediction(prompt, client):
    response = client.complete(
        messages=[
            SystemMessage(content="You are a helpful assistant."),
            UserMessage(content=prompt),
        ]
    )
    
    try:
        return parse_response(response)
    except Exception as e:
        print(f"Failed to get response: {e}")
        return None

def parse_response(response):
    if response.choices:
        choice = response.choices[0]
        if choice.message and choice.message.content:
            content = choice.message.content
            print(f"Model response content: {content}")
            return content
        if choice.finish_reason:
            finish_reason = choice.finish_reason
            print(f"Finish reason: {finish_reason}")
    return "No valid response"

if __name__ == "__main__":
    if not MODEL_1_ENDPOINT_URL or not MODEL_1_PRIMARY_KEY or not MODEL_2_ENDPOINT_URL or not MODEL_2_PRIMARY_KEY:
        print("Error: Missing endpoint URL or primary key for one or both models")
    else:
        print("Response from Model 1:")
        response1 = get_model_prediction(PROMPT, client1)
        if response1:
            print(response1)
        else:
            print("No response or failed to decode the response from Model 1.")

        print("\nResponse from Model 2:")
        response2 = get_model_prediction(PROMPT, client2)
        if response2:
            print(response2)
        else:
            print("No response or failed to decode the response from Model 2.")


[
    {"question": "What is the largest desert in the world?", "expected": "Sahara"},
    {"question": "Which river is the longest in the world?", "expected": "Nile"},
    {"question": "What is the capital of Japan?", "expected": "Tokyo"},
    {"question": "Which ocean is the largest by surface area?", "expected": "Pacific Ocean"},
    {"question": "Which desert is known as the 'Cold Desert'?", "expected": "Gobi Desert"}
]


import requests
import json
import time
import asyncio
import weave
from weave import Evaluation, Model
from azure.ai.inference import ChatCompletionsClient
from azure.ai.inference.models import SystemMessage, UserMessage
from azure.core.credentials import AzureKeyCredential


# Initialize Weave for logging
weave.init('azure-api-eval-geo')


# Define global variables for the two models
MODEL_1_ENDPOINT_URL = "https://Mistral-small-achqr.eastus2.models.ai.azure.com/chat/completions"
MODEL_1_PRIMARY_KEY = "your key"

MODEL_2_ENDPOINT_URL = "https://Phi-3-small-128k-instruct-fcsez.eastus2.models.ai.azure.com/chat/completions"
MODEL_2_PRIMARY_KEY = "your key"


# Initialize the ChatCompletionsClient for each model
def init_client(endpoint, key):
    return ChatCompletionsClient(
        endpoint=endpoint,
        credential=AzureKeyCredential(key)
    )


def get_model_prediction(prompt, client, model_name):
    response = client.complete(
        messages=[
            SystemMessage(content="You are a helpful assistant."),
            UserMessage(content=prompt),
        ]
    )
    try:
        content = response.choices[0].message.content
        prompt_tokens = response.usage.prompt_tokens
        completion_tokens = response.usage.completion_tokens
        total_tokens = response.usage.total_tokens
        
        # Update the returned result to include the correct token structure
        return {
            'output': content,
            'model': model_name,            
            'usage': {
                'input_tokens': prompt_tokens,
                'output_tokens': completion_tokens,
                'total_tokens' : total_tokens
            }
        }
    except Exception as e:
        print(f"Failed to get response: {e}")
        return {
            'output': "No valid response",
            'usage': {
                'input_tokens': 0,
                'output_tokens': 0,
                'total_tokens': 0
            }
        }

# Define custom scoring functions
@weave.op()
def substring_match(expected: str, model_output: dict) -> dict:
    match = expected.lower() in model_output['output'].lower()
    return {
        'substring_match': match
    }


# Define the base model for Weave
class MistralSmall(Model):
    model_name: str = "Mistral-small"
    endpoint: str = MODEL_1_ENDPOINT_URL
    primary_key: str = MODEL_1_PRIMARY_KEY
    client: ChatCompletionsClient = None  # Declare the client field

    def __init__(self):
        super().__init__()  # Initialize the base class if needed
        # Initialize the client and assign it explicitly
        object.__setattr__(self, 'client', init_client(self.endpoint, self.primary_key))

    @weave.op()
    def predict(self, question: str):
        result = get_model_prediction(question, self.client, self.model_name)
        return result


class Phi3Small(Model):
    model_name: str = "Phi-3-small"
    endpoint: str = MODEL_2_ENDPOINT_URL
    primary_key: str = MODEL_2_PRIMARY_KEY
    client: ChatCompletionsClient = None  # Declare the client field

    def __init__(self):
        super().__init__()  # Initialize the base class if needed
        # Initialize the client and assign it explicitly
        object.__setattr__(self, 'client', init_client(self.endpoint, self.primary_key))
    

    @weave.op()
    def predict(self, question: str):
        result = get_model_prediction(question, self.client, self.model_name)
        return result

def run_evaluation(model_class):
    # Load the dataset
    with open('dataset.json', 'r') as f:
        dataset = json.load(f)

    # Instantiate the model
    model = model_class()

    # Create evaluation for the model
    evaluation = Evaluation(dataset=dataset, scorers=[substring_match])

    # Run the evaluation for the model
    print(f"Evaluating {model.model_name}:")
    asyncio.run(evaluation.evaluate(model.predict))
    print(f"{model.model_name} evaluation completed.")


if __name__ == "__main__":
    # Run evaluation for both models sequentially
    run_evaluation(MistralSmall)
    run_evaluation(Phi3Small)