|
| 1 | +#%% |
| 2 | + |
| 3 | +from twinkle_client import init_tinker_compat_client |
| 4 | +service_client = init_tinker_compat_client(base_url='http://localhost:8000') |
| 5 | + |
| 6 | +print("Available models:") |
| 7 | +for item in service_client.get_server_capabilities().supported_models: |
| 8 | + print("- " + item.model_name) |
| 9 | + |
| 10 | +#%% |
| 11 | +base_model = "ms://Qwen/Qwen2.5-0.5B-Instruct" |
| 12 | +training_client = service_client.create_lora_training_client( |
| 13 | + base_model=base_model |
| 14 | +) |
| 15 | + |
| 16 | +#%% |
| 17 | +# Create some training examples |
| 18 | +examples = [ |
| 19 | + {"input": "banana split", "output": "anana-bay plit-say"}, |
| 20 | + {"input": "quantum physics", "output": "uantum-qay ysics-phay"}, |
| 21 | + {"input": "donut shop", "output": "onut-day op-shay"}, |
| 22 | + {"input": "pickle jar", "output": "ickle-pay ar-jay"}, |
| 23 | + {"input": "space exploration", "output": "ace-spay exploration-way"}, |
| 24 | + {"input": "rubber duck", "output": "ubber-ray uck-day"}, |
| 25 | + {"input": "coding wizard", "output": "oding-cay izard-way"}, |
| 26 | +] |
| 27 | + |
| 28 | +# Convert examples into the format expected by the training client |
| 29 | +from tinker import types |
| 30 | +from modelscope import AutoTokenizer |
| 31 | +# Get the tokenizer from the training client |
| 32 | +# tokenizer = training_client.get_tokenizer() # NOTE: network call huggingface |
| 33 | +tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", trust_remote_code=True) |
| 34 | + |
| 35 | +def process_example(example: dict, tokenizer) -> types.Datum: |
| 36 | + # Format the input with Input/Output template |
| 37 | + # For most real use cases, you'll want to use a renderer / chat template, |
| 38 | + # (see later docs) but here, we'll keep it simple. |
| 39 | + prompt = f"English: {example['input']}\nPig Latin:" |
| 40 | + |
| 41 | + prompt_tokens = tokenizer.encode(prompt, add_special_tokens=True) |
| 42 | + prompt_weights = [0] * len(prompt_tokens) |
| 43 | + # Add a space before the output string, and finish with double newline |
| 44 | + completion_tokens = tokenizer.encode(f" {example['output']}\n\n", add_special_tokens=False) |
| 45 | + completion_weights = [1] * len(completion_tokens) |
| 46 | + |
| 47 | + tokens = prompt_tokens + completion_tokens |
| 48 | + weights = prompt_weights + completion_weights |
| 49 | + |
| 50 | + input_tokens = tokens[:-1] |
| 51 | + target_tokens = tokens[1:] # We're predicting the next token, so targets need to be shifted. |
| 52 | + weights = weights[1:] |
| 53 | + |
| 54 | + # A datum is a single training example for the loss function. |
| 55 | + # It has model_input, which is the input sequence that'll be passed into the LLM, |
| 56 | + # loss_fn_inputs, which is a dictionary of extra inputs used by the loss function. |
| 57 | + return types.Datum( |
| 58 | + model_input=types.ModelInput.from_ints(tokens=input_tokens), |
| 59 | + loss_fn_inputs=dict(weights=weights, target_tokens=target_tokens) |
| 60 | + ) |
| 61 | + |
| 62 | +processed_examples = [process_example(ex, tokenizer) for ex in examples] |
| 63 | + |
| 64 | +# Visualize the first example for debugging purposes |
| 65 | +datum0 = processed_examples[0] |
| 66 | +print(f"{'Input':<20} {'Target':<20} {'Weight':<10}") |
| 67 | +print("-" * 50) |
| 68 | +for i, (inp, tgt, wgt) in enumerate(zip(datum0.model_input.to_ints(), datum0.loss_fn_inputs['target_tokens'].tolist(), datum0.loss_fn_inputs['weights'].tolist())): |
| 69 | + print(f"{repr(tokenizer.decode([inp])):<20} {repr(tokenizer.decode([tgt])):<20} {wgt:<10}") |
| 70 | + |
| 71 | +#%% |
| 72 | +import numpy as np |
| 73 | +for _ in range(6): |
| 74 | + fwdbwd_future = training_client.forward_backward(processed_examples, "cross_entropy") |
| 75 | + optim_future = training_client.optim_step(types.AdamParams(learning_rate=1e-4)) |
| 76 | + |
| 77 | + # Wait for the results |
| 78 | + fwdbwd_result = fwdbwd_future.result() |
| 79 | + optim_result = optim_future.result() |
| 80 | + |
| 81 | + # fwdbwd_result contains the logprobs of all the tokens we put in. Now we can compute the weighted |
| 82 | + # average log loss per token. |
| 83 | + logprobs = np.concatenate([output['logprobs'].tolist() for output in fwdbwd_result.loss_fn_outputs]) |
| 84 | + weights = np.concatenate([example.loss_fn_inputs['weights'].tolist() for example in processed_examples]) |
| 85 | + print(f"Loss per token: {-np.dot(logprobs, weights) / weights.sum():.4f}") |
| 86 | + |
| 87 | +#%% |
| 88 | +# First, create a sampling client. We need to transfer weights |
| 89 | +sampling_client = training_client.save_weights_and_get_sampling_client(name='pig-latin-model') |
| 90 | + |
| 91 | +# Now, we can sample from the model. |
| 92 | +prompt = types.ModelInput.from_ints(tokenizer.encode("English: coffee break\nPig Latin:")) |
| 93 | +params = types.SamplingParams(max_tokens=20, temperature=0.0, stop=["\n"]) # Greedy sampling |
| 94 | +future = sampling_client.sample(prompt=prompt, sampling_params=params, num_samples=8) |
| 95 | +result = future.result() |
| 96 | +print("Responses:") |
| 97 | +for i, seq in enumerate(result.sequences): |
| 98 | + print(f"{i}: {repr(tokenizer.decode(seq.tokens))}") |
| 99 | +# %% |
0 commit comments