Add prompt tuning to the basic example

[justheuristic]

The following code can be added to basic example to run prompt tuning and show demonstrate that it works

Quest: go to the main example notebook (in readme), and add some simple example based on prompt tuning

Here's a basic example. I'm not sure it aint missing something, but if it is, we can help add it back.

import torch
from transformers import BloomTokenizerFast
from petals.client import DistributedBloomForCausalLM
assert 'model' not in globals(), "please restart the kernel"
 
MODEL_NAME = "bigscience/bloom-petals"
tokenizer = BloomTokenizerFast.from_pretrained(MODEL_NAME, padding_side='right')
model = DistributedBloomForCausalLM.from_pretrained(
    MODEL_NAME, tuning_mode='deep_ptune', pre_seq_len=3
).cuda()

inputs = tokenizer("A quick brown fox ", return_tensors="pt")["input_ids"].cuda()
remote_outputs = model.generate(inputs, max_new_tokens=7)
print("generated:", tokenizer.decode(remote_outputs[0]))
 
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

the_fox_is_innocent = tokenizer("A quick brown fox did not jump over the lazy dog", return_tensors="pt")["input_ids"].cuda()
for i in range(50):
  loss = model(input_ids=the_fox_is_innocent, labels=the_fox_is_innocent).loss
  opt.zero_grad()
  loss.backward()
  opt.step()
  print(f"loss[{i}] = {loss.item():.3f}")

inputs = tokenizer("A quick brown fox ", return_tensors="pt")["input_ids"].cuda()
remote_outputs = model.generate(inputs, max_new_tokens=7)
print("generated:", tokenizer.decode(remote_outputs[0]))

[borzunov]

Done before the 2.0.0 release.