Update app.py

app.py

@@ -8,24 +8,15 @@ from src.client import DistributedBloomForCausalLM
 INITIAL_PEERS = ['/ip6/2a0b:4880::a242:3fff:fe3a:2ae1/tcp/21338/p2p/QmSXDXLeSMXjS4YerDrdn1zpGQaNzkZ9ogN2SoAEyAdDhs', '/ip6/2a0b:4880::a242:3fff:fe3a:2ae1/udp/21338/quic/p2p/QmSXDXLeSMXjS4YerDrdn1zpGQaNzkZ9ogN2SoAEyAdDhs']
 
 tokenizer = transformers.BloomTokenizerFast.from_pretrained("bigscience/test-bloomd-6b3")
-model = DistributedBloomForCausalLM.from_pretrained("bigscience/test-bloomd-6b3", initial_peers=INITIAL_PEERS, low_cpu_mem_usage=True, torch_dtype=torch.float32)
+#model = DistributedBloomForCausalLM.from_pretrained("bigscience/test-bloomd-6b3", initial_peers=INITIAL_PEERS, low_cpu_mem_usage=True, torch_dtype=torch.float32)
 
 def inference(text, seq_length=1):
-    input_ids = tokenizer(text, return_tensors='pt')['input_ids']
-    with torch.inference_mode(), model.transformer.h.inference_session() as remote_transformer:
-        for i in range(seq_length):
-            h = model.transformer.word_embeddings(input_ids)
-            h = model.transformer.word_embeddings_layernorm(h)
-
-            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
-
-            h = model.transformer.ln_f(h)
-            h = F.linear(h, weight=model.transformer.word_embeddings.weight)  # note: this line takes a while, will also be fixed
-            next_token_ix = torch.multinomial((h[0, -1] / 0.8).softmax(-1), 1)
-            
-            # print(end=tokenizer.decode(next_token_ix.item()))
-            input_ids = next_token_ix.view(1, 1)
-    return tokenizer.decode(input_ids.item())
+    #input_ids = tokenizer(text, return_tensors='pt')['input_ids']
+    #with torch.inference_mode(), model.transformer.h.inference_session() as remote_transformer:
+    #    for i in range(seq_length):
+    #        h = model.transformer.word_embeddings(input_ids)
+    #        h = model.transformer.word_embeddings_layernorm(h)
+    return text[::-1]
 
 iface = gr.Interface(fn=inference, inputs="text", outputs="text")
 iface.launch()