Younes Belkada commited on
Commit
da8e0b6
·
1 Parent(s): 7e52371

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -16
app.py CHANGED
@@ -5,11 +5,13 @@ import gradio as gr
5
 
6
  from src.client import DistributedBloomForCausalLM
7
 
8
- INITIAL_PEERS = [
9
- '/ip4/193.106.95.184/tcp/31337/p2p/QmUigSxrVz9x5FR9ZYr4iRfEX2vDxihL2YZtDd7sp2eKnM',
10
- '/ip6/193.106.95.184/tcp/21337/p2p/QmSXDXLeSMXjS4YerDrdn1zpGQaNzkZ9ogN2SoAEyAdDhs',
11
- '/ip6/193.106.95.184/udp/21337/quic/QmSXDXLeSMXjS4YerDrdn1zpGQaNzkZ9ogN2SoAEyAdDhs',
12
- ]
 
 
13
  tokenizer = transformers.BloomTokenizerFast.from_pretrained("bigscience/test-bloomd-6b3")
14
  model = DistributedBloomForCausalLM.from_pretrained("bigscience/test-bloomd-6b3", initial_peers=INITIAL_PEERS, low_cpu_mem_usage=True, torch_dtype=torch.float32)
15
 
@@ -19,16 +21,8 @@ def inference(text, seq_length=1):
19
  for i in range(seq_length):
20
  h = model.transformer.word_embeddings(input_ids)
21
  h = model.transformer.word_embeddings_layernorm(h)
22
-
23
- h = remote_transformer.step(h) # note [yozh]: this line currently freezes for 10 seconds first time only, its gonna be fixed in the nearest PR
24
-
25
- h = model.transformer.ln_f(h)
26
- h = F.linear(h, weight=model.transformer.word_embeddings.weight) # note: this line takes a while, will also be fixed
27
- next_token_ix = torch.multinomial((h[0, -1] / 0.8).softmax(-1), 1)
28
-
29
- # print(end=tokenizer.decode(next_token_ix.item()))
30
- input_ids = next_token_ix.view(1, 1)
31
- return tokenizer.decode(input_ids.item())
32
-
33
  iface = gr.Interface(fn=inference, inputs="text", outputs="text")
34
  iface.launch()
 
5
 
6
  from src.client import DistributedBloomForCausalLM
7
 
8
+ INITIAL_PEERS = ['/ip4/193.106.95.184/tcp/443/p2p/QmSXDXLeSMXjS4YerDrdn1zpGQaNzkZ9ogN2SoAEyAdDhs']
9
+
10
+ import hivemind # test that DHT instances work on localhost
11
+ dht1 = hivemind.DHT(start=True)
12
+ dht2 = hivemind.DHT(start=True, initial_peers=dht1.get_visible_maddrs())
13
+
14
+
15
  tokenizer = transformers.BloomTokenizerFast.from_pretrained("bigscience/test-bloomd-6b3")
16
  model = DistributedBloomForCausalLM.from_pretrained("bigscience/test-bloomd-6b3", initial_peers=INITIAL_PEERS, low_cpu_mem_usage=True, torch_dtype=torch.float32)
17
 
 
21
  for i in range(seq_length):
22
  h = model.transformer.word_embeddings(input_ids)
23
  h = model.transformer.word_embeddings_layernorm(h)
24
+ h = remote_transformer.step(h)
25
+ return repr(h)
26
+
 
 
 
 
 
 
 
 
27
  iface = gr.Interface(fn=inference, inputs="text", outputs="text")
28
  iface.launch()