1 line
10 KiB
JSON
1 line
10 KiB
JSON
{"repo": "aprxi/talu", "files": {"examples/developers/chat/production_features.py": {"c850924a91aaa00a679ba07bb597b84500ccfe4d": {"t": 1785674491.409504, "n": 9595, "h": "45ac409d80b83617", "c": "\"\"\"Production Features - count_tokens, logit_bias, extra_body.\n\nPrimary API: talu.Chat, talu.GenerationConfig\nScope: Single\n\nThis example demonstrates production-ready features that enable real-world\ndeployment scenarios like RAG pipelines, classification, and provider integration.\n\nKey points:\n- count_tokens() provides accurate token counts including template overhead\n- logit_bias enables controlled generation for classification\n- extra_body is the \"escape hatch\" for provider-specific parameters\n\nRelated:\n - examples/developers/chat/configuration.py (basic config)\n - examples/developers/chat/config_composition.py (merging configs)\n\"\"\"\n\nfrom talu import Chat\nfrom talu.router import GenerationConfig\n\n# =============================================================================\n# 1. Context Window Management (count_tokens)\n# =============================================================================\n\nprint(\"=== Context Window Management ===\")\n\nchat = Chat(\"Qwen/Qwen3-0.6B\", system=\"You are a helpful assistant.\")\n\n# Count current tokens (includes system prompt + template overhead)\ncurrent_tokens = chat.count_tokens()\nprint(f\"Base context (system + template): {current_tokens} tokens\")\n\n# Simulate adding a document\ndocument = \"This is a long document about machine learning. \" * 50\n\n# Check if document fits before adding\nif chat.count_tokens(document) < 2048: # Assume 2048 context limit\n print(f\"Document would use: {chat.count_tokens(document)} tokens\")\n print(\"Document fits in context!\")\nelse:\n print(\"Document too large - need to summarize or chunk\")\n\n# After sending a message, count updates\nresponse = chat.send(\"What is 2+2?\", max_tokens=10)\nprint(f\"After one exchange: {chat.count_tokens()} tokens\")\n\n\n# =============================================================================\n# 2. Controlled Generation (logit_bias)\n# =============================================================================\n\nprint(\"\\n=== Controlled Generation (logit_bias) ===\")\n\n# For classification, force the model to output specific tokens\n# First, get token IDs for your target words using the tokenizer:\n# tokenizer = Tokenizer(\"model\")\n# yes_id = tokenizer.encode(\"Yes\")[0]\n# no_id = tokenizer.encode(\"No\")[0]\n\n# Example token IDs (these vary by model)\nYES_TOKEN_ID = 9891\nNO_TOKEN_ID = 2841\n\nclassification_config = GenerationConfig(\n temperature=0.0, # Deterministic\n max_tokens=1, # Only need one token\n logit_bias={\n YES_TOKEN_ID: 100.0, # Strongly prefer\n NO_TOKEN_ID: 100.0, # Strongly prefer\n # All other tokens implicitly get bias 0\n },\n)\n\nprint(\"Classification config with logit_bias:\")\nprint(f\" temperature: {classification_config.temperature}\")\nprint(f\" max_tokens: {classification_config.max_tokens}\")\nprint(f\" logit_bias: {classification_config.logit_bias}\")\n\n# Use for spam classification:\n# chat = Chat(\"model\")\n# response = chat.send(\n# f\"Is this spam? Answer Yes or No only.\\n\\n{email_text}\",\n# config=classification_config\n# )\n# is_spam = \"Yes\" in str(response)\n\n\n# =============================================================================\n# 3. Provider-Specific Parameters (extra_body)\n# =============================================================================\n\nprint(\"\\n=== Provider-Specific Parameters (extra_body) ===\")\n\n# extra_body passes arbitrary JSON to remote APIs\n# Useful for bleeding-edge parameters not yet in GenerationConfig\n\n# vLLM-specific parameters\nvllm_config = GenerationConfig(\n temperature=0.7,\n max_tokens=100,\n extra_body={\n \"best_of\": 3, # vLLM: generate 3, return best\n \"use_beam_search\": False, # vLLM: disable beam search\n \"presence_penalty\": 0.5, # OpenAI-compatible\n \"frequency_penalty\": 0.3, # OpenAI-compatible\n },\n)\n\nprint(\"vLLM config with extra_body:\")\nprint(f\" standard: temp={vllm_config.temperature}, max_tokens={vllm_config.max_tokens}\")\nprint(f\" extra_body: {vllm_config.extra_body}\")\n\n# Together.ai-specific parameters\ntogether_config = GenerationConfig(\n temperature=0.7,\n extra_body={\n \"repetition_penalty\": 1.1, # Together-specific\n \"top_k_return_sequences\": 1, # Together-specific\n },\n)\n\nprint(\"\\nTogether.ai config with extra_body:\")\nprint(f\" extra_body: {together_config.extra_body}\")\n\n# Ollama-specific parameters\nollama_config = GenerationConfig(\n temperature=0.7,\n extra_body={\n \"num_ctx\": 4096, # Ollama: context window\n \"num_gpu\": 1, # Ollama: GPU layers\n \"mirostat\": 2, # Ollama: Mirostat sampling\n \"mirostat_eta\": 0.1,\n \"mirostat_tau\": 5.0,\n },\n)\n\nprint(\"\\nOllama config with extra_body:\")\nprint(f\" extra_body: {ollama_config.extra_body}\")\n\n\n# =============================================================================\n# 4. RAG Pipeline Pattern\n# =============================================================================\n\nprint(\"\\n=== RAG Pipeline Pattern ===\")\n\n# RAG budgeting: ensure context fits before generation\n\ndef rag_generate(chat: Chat, context: str, query: str, max_context: int = 4096):\n \"\"\"Generate with RAG context, respecting context limits.\"\"\"\n # Reserve tokens for response\n response_budget = 512\n available = max_context - response_budget\n\n # Check if context fits\n prompt = f\"Context:\\n{context}\\n\\nQuestion: {query}\"\n needed = chat.count_tokens(prompt)\n\n if needed > available:\n # Truncate context (in practice: chunk or summarize)\n print(f\"Context too large ({needed} > {available}), truncating...\")\n # Simple truncation - real apps would chunk/summarize\n context = context[: len(context) // 2]\n prompt = f\"Context:\\n{context}\\n\\nQuestion: {query}\"\n\n print(f\"Final prompt: {chat.count_tokens(prompt)} tokens\")\n return chat.send(prompt, max_tokens=response_budget)\n\n\n# Example usage:\n# context = fetch_documents(query)\n# response = rag_generate(chat, context, \"What is the main topic?\")\n\n\n# =============================================================================\n# 5. A/B Testing with Different Providers\n# =============================================================================\n\nprint(\"\\n=== A/B Testing Pattern ===\")\n\n# Base config for both providers\nbase_config = GenerationConfig(\n temperature=0.7,\n max_tokens=200,\n stop_sequences=[\"\\n\\n\"],\n)\n\n# Provider A: vLLM with specific tuning\nprovider_a = base_config | GenerationConfig(\n extra_body={\"best_of\": 2, \"presence_penalty\": 0.5}\n)\n\n# Provider B: Together with different tuning\nprovider_b = base_config | GenerationConfig(\n extra_body={\"repetition_penalty\": 1.1}\n)\n\nprint(\"A/B test configs:\")\nprint(f\" Provider A (vLLM): {provider_a.extra_body}\")\nprint(f\" Provider B (Together): {provider_b.extra_body}\")\n\n# In production:\n# if random.random() < 0.5:\n# response = chat_vllm.send(prompt, config=provider_a)\n# else:\n# response = chat_together.send(prompt, config=provider_b)\n\n\n# =============================================================================\n# 6. Custom Headers for Enterprise Networking\n# =============================================================================\n\nprint(\"\\n=== Custom Headers for Enterprise ===\")\n\nfrom talu.router import ModelSpec, OpenAICompatibleBackend\n\n# Enterprise pattern 1: Internal proxy with custom auth\n# Many enterprises route LLM traffic through internal proxies that require\n# custom headers for authentication, routing, or tracing.\n\ninternal_proxy_backend = OpenAICompatibleBackend(\n base_url=\"https://internal-llm-proxy.corp.example.com/v1\",\n api_key=None, # Proxy handles authentication\n headers={\n \"X-Proxy-Auth\": \"internal-service-token\",\n \"X-Service-Name\": \"ml-pipeline\",\n \"X-Team-ID\": \"data-science\",\n \"X-Correlation-ID\": \"request-12345\", # For distributed tracing\n },\n)\n\nprint(\"Internal proxy config:\")\nprint(f\" base_url: {internal_proxy_backend.base_url}\")\nprint(f\" headers: {internal_proxy_backend.headers}\")\n\n# Enterprise pattern 2: OpenAI with request tracing\n# Add custom headers alongside the standard OpenAI authentication.\n\nopenai_with_tracing = OpenAICompatibleBackend(\n base_url=\"https://api.openai.com/v1\",\n api_key=\"sk-...\", # Your OpenAI key\n org_id=\"org-123\",\n headers={\n \"X-Request-ID\": \"trace-abc-123\",\n \"X-Environment\": \"production\",\n },\n)\n\nprint(\"\\nOpenAI with tracing:\")\nprint(f\" api_key: {openai_with_tracing.api_key[:10]}...\")\nprint(f\" org_id: {openai_with_tracing.org_id}\")\nprint(f\" headers: {openai_with_tracing.headers}\")\n\n# Enterprise pattern 3: AWS Bedrock-style auth via SigV4 headers\n# (Hypothetical - actual SigV4 signing would be done separately)\n\nbedrock_style = OpenAICompatibleBackend(\n base_url=\"https://bedrock.us-east-1.amazonaws.com/v1\",\n headers={\n \"X-Amz-Target\": \"AmazonBedrock.Converse\",\n \"X-Amz-Content-Sha256\": \"...\",\n \"Authorization\": \"AWS4-HMAC-SHA256 Credential=...\",\n },\n)\n\nprint(\"\\nBedrock-style config:\")\nprint(f\" headers: {list(bedrock_style.headers.keys())}\")\n\n# Using headers with Chat\n# spec = ModelSpec(\n# ref=\"gpt-4o\",\n# backend=internal_proxy_backend,\n# )\n# chat = Chat(spec, system=\"You are a helpful assistant.\")\n# response = chat.send(\"Hello!\")\n\n\n\"\"\"\nTopics covered:\n* generation.config\n* chat.context\n* config.sampling\n* client.ask\n* chat.streaming\n* workflow.end.to.end\n\nRelated:\n* examples/developers/chat/configuration.py\n* examples/developers/chat/config_composition.py\n* examples/developers/chat/custom_endpoints.py\n\"\"\"\n"}}}} |