File size: 1,586 Bytes
9644d0b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

import asyncio
import torch
import os
import sys
from transformers import AutoModelForCausalLM, AutoTokenizer

# Ensure we can import from all_space
sys.path.insert(0, os.getcwd())

from all_space.generators import generate_chat_completion
from all_space.px_patches.gemma3_270m_px_baseline.patch import apply_px_patch

async def run_official_style_test(model_id="google/gemma-3-270m-it", config_preset="SUBJECTIVE"):
    print(f"--- Official Style Coherence Test: {model_id} (Preset={config_preset}) ---")
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="cuda")
    
    if config_preset != "BASELINE":
        apply_px_patch(model, config_preset=config_preset)
    
    model_entry = {
        "model": model,
        "tokenizer": tokenizer
    }
    
    test_prompts = [
        "What is the capital of France?",
        "Solve: 15 + 27 * 2"
    ]
    
    for prompt in test_prompts:
        print(f"\nPrompt: {prompt}")
        messages = [{"role": "user", "content": prompt}]
        
        result = await generate_chat_completion(
            model_entry=model_entry,
            messages=messages,
            temperature=0.7,
            top_p=0.9,
            max_tokens=128
        )
        
        print(f"Response: {result['text']}")

if __name__ == "__main__":
    asyncio.run(run_official_style_test(config_preset="BASELINE"))
    asyncio.run(run_official_style_test(config_preset="SUBJECTIVE"))
    asyncio.run(run_official_style_test(config_preset="RIGOR"))