Pythonollama_client.py
Copy
1
import time
2
import openai  # pip install openai
3
 
4
# Ollama exposes OpenAI-compatible API on localhost:11434
5
# Run: ollama pull llama3.2:3b && ollama serve
6
 
7
client = openai.OpenAI(
8
    base_url="http://localhost:11434/v1",
9
    api_key="ollama",  # Required but ignored
10
)
11
 
12
def benchmark_model(model: str, prompt: str, n: int = 3) -> dict:
13
    """Benchmark tokens/sec for a local model."""
14
    latencies = []
15
    token_counts = []
16
    
17
    for _ in range(n):
18
        start = time.perf_counter()
19
        response = client.chat.completions.create(
20
            model=model,
21
            messages=[{"role": "user", "content": prompt}],
22
            max_tokens=200,
23
            stream=False,
24
        )
25
        latency = time.perf_counter() - start
26
        latencies.append(latency)
27
        token_counts.append(response.usage.completion_tokens)
28
    
29
    avg_latency = sum(latencies) / n
30
    avg_tokens = sum(token_counts) / n
31
    tokens_per_sec = avg_tokens / avg_latency
32
    
33
    return {
34
        "model": model,
35
        "avg_latency_s": round(avg_latency, 2),
36
        "avg_output_tokens": round(avg_tokens),
37
        "tokens_per_sec": round(tokens_per_sec, 1),
38
    }
39
 
40
# ── Streaming example ───────────────────────────────────────────────
41
def stream_local(model: str, prompt: str):
42
    """Stream tokens from local model."""
43
    print(f"\n[{model}] Streaming: ", end="")
44
    start = time.perf_counter()
45
    token_count = 0
46
    
47
    stream = client.chat.completions.create(
48
        model=model,
49
        messages=[{"role": "user", "content": prompt}],
50
        max_tokens=200,
51
        stream=True,
52
    )
53
    
54
    for chunk in stream:
55
        if chunk.choices[0].delta.content:
56
            print(chunk.choices[0].delta.content, end="", flush=True)
57
            token_count += 1
58
    
59
    elapsed = time.perf_counter() - start
60
    print(f"\n  Tokens/sec: {token_count/elapsed:.1f}")
61
 
62
PROMPT = "Explain gradient descent in 3 sentences."
63
 
64
# Benchmark available models
65
for model in ["llama3.2:3b"]:
66
    try:
67
        result = benchmark_model(model, PROMPT)
68
        print(result)
69
        stream_local(model, PROMPT)
70
    except Exception as e:
71
        print(f"{model}: {e} (model not pulled — run: ollama pull {model})")
72