d84f98799b
## Added ### Core framework - src/llm_client.py — Ollama API wrapper for local Docker inference - src/samples.py — Embedded benchmark datasets (GSM8K/MATH/AIME + Chinese) - src/run_all.py — Unified experiment runner ### Phase 1: Baselines - src/baseline/benchmark.py — IO, CoT, CoT-SC, ToT evaluation with Chinese support ### Phase 2: AGoT - src/agot/agot.py — AGoT core algorithm (6 agent types, recursive decomposition) - src/agot/graph_utils.py — DAG graph data structure with cycle detection - src/agot/prompts.py — 6 agent prompt templates (EN + ZH) - src/agot/run.py — AGoT experiment runner ### Graph Analysis - src/graph_analysis/reasoning_graph.py — Graph property computation (cyclicity, diameter, small-world) - src/graph_analysis/visualize.py — Visualization charts ### Documentation - docs/graph_reasoning_principles.md — Comprehensive principles document - docs/next_steps.md — Roadmap for reliable reasoning tool ### Experiment Results - GSM8K/MATH baselines on qwen2.5:32b and qwq:latest - AGoT validation (100% on small sample) - Chinese dataset benchmarks (C-GSM8K, CMATH) - Graph property analysis confirming Topology of Reasoning findings
114 lines
3.9 KiB
Python
114 lines
3.9 KiB
Python
"""
|
|
Ollama LLM Client — wraps the Ollama REST API running in Docker.
|
|
|
|
Provides a unified interface for all experiment scripts to call local models.
|
|
"""
|
|
|
|
import json
|
|
import time
|
|
from typing import Optional
|
|
import ollama
|
|
|
|
# ── Default Configuration ─────────────────────────────────────────────
|
|
DEFAULT_BASE_URL = "http://localhost:11434"
|
|
DEFAULT_MODEL = "qwen2.5:32b"
|
|
|
|
|
|
class LLMClient:
|
|
"""A lightweight wrapper around the Ollama API."""
|
|
|
|
def __init__(
|
|
self,
|
|
model: str = DEFAULT_MODEL,
|
|
base_url: str = DEFAULT_BASE_URL,
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
max_tokens: int = 2048,
|
|
num_ctx: int = 8192,
|
|
seed: Optional[int] = None,
|
|
):
|
|
self.model = model
|
|
self.client = ollama.Client(host=base_url)
|
|
self.options = dict(
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
num_predict=max_tokens,
|
|
num_ctx=num_ctx,
|
|
)
|
|
if seed is not None:
|
|
self.options["seed"] = seed
|
|
|
|
# ── Public API ────────────────────────────────────────────────────
|
|
|
|
def generate(self, prompt: str, system: Optional[str] = None) -> str:
|
|
"""Simple text generation (no chat history)."""
|
|
messages = []
|
|
if system:
|
|
messages.append({"role": "system", "content": system})
|
|
messages.append({"role": "user", "content": prompt})
|
|
return self._chat(messages)
|
|
|
|
def chat(self, messages: list) -> str:
|
|
"""Multi-turn chat. messages = [{"role": "user"/"assistant", "content": ...}]."""
|
|
return self._chat(messages)
|
|
|
|
def generate_with_metadata(self, prompt: str, system: Optional[str] = None) -> dict:
|
|
"""Like generate() but returns full response metadata."""
|
|
messages = []
|
|
if system:
|
|
messages.append({"role": "system", "content": system})
|
|
messages.append({"role": "user", "content": prompt})
|
|
return self._chat_full(messages)
|
|
|
|
def count_tokens(self, text: str) -> int:
|
|
"""Estimate token count (rough: ~4 chars per token for Chinese+English)."""
|
|
return len(text) // 3 + 1
|
|
|
|
@property
|
|
def model_name(self) -> str:
|
|
return self.model
|
|
|
|
# ── Internal ──────────────────────────────────────────────────────
|
|
|
|
def _chat(self, messages: list) -> str:
|
|
resp = self._chat_full(messages)
|
|
return resp["message"]["content"]
|
|
|
|
def _chat_full(self, messages: list) -> dict:
|
|
resp = self.client.chat(
|
|
model=self.model,
|
|
messages=messages,
|
|
options=self.options,
|
|
)
|
|
return resp
|
|
|
|
def __repr__(self) -> str:
|
|
return f"LLMClient(model={self.model})"
|
|
|
|
|
|
# ── Convenience Factory ───────────────────────────────────────────────
|
|
|
|
def make_client(
|
|
model: str = DEFAULT_MODEL,
|
|
temperature: float = 0.0,
|
|
max_tokens: int = 2048,
|
|
seed: int = 42,
|
|
) -> LLMClient:
|
|
"""Create a deterministic client (temperature=0) suitable for evaluation."""
|
|
return LLMClient(
|
|
model=model,
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
seed=seed,
|
|
)
|
|
|
|
|
|
# ── Quick Test ────────────────────────────────────────────────────────
|
|
|
|
if __name__ == "__main__":
|
|
# Test the client
|
|
client = make_client()
|
|
print(f"Model: {client.model_name}")
|
|
resp = client.generate("Say hello in exactly one word.")
|
|
print(f"Response: {resp}")
|