Update app.py

app.py

@@ -17,12 +17,19 @@ model = DistributedBloomForCausalLM.from_pretrained("bigscience/test-bloomd-6b3"
 
 def inference(text, seq_length=1):
     input_ids = tokenizer(text, return_tensors='pt')['input_ids']
+    final_tokens = 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)
-            return repr(h)
-    
+            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)
+        
+            final_tokens = torch.cat([final_tokens, next_token_ix], dim=-1)
+            input_ids = next_token_ix.view(1, 1)
+         return tokenizer.decode(final_tokens, skip_special_tokens=False)
+ 
 iface = gr.Interface(fn=inference, inputs="text", outputs="text")
 iface.launch()