Simplify usage

README.md

@@ -1,5 +1,9 @@
 ---
 pipeline_tag: text-generation
+base_model:
+- google/gemma-3-1b-it
+library_name: transformers.js
+license: gemma
 ---
 
 ## Usage
@@ -41,8 +45,7 @@ past_key_values = {
     for kv in ('key', 'value')
 }
 input_ids = inputs['input_ids']
-attention_mask = inputs['attention_mask']
-position_ids = np.cumsum(inputs['attention_mask'], axis=-1)
+position_ids = np.tile(np.arange(1, input_ids.shape[-1] + 1), (batch_size, 1))
 
 # 3. Generation loop
 max_new_tokens = 1024
@@ -50,14 +53,12 @@ generated_tokens = np.array([[]], dtype=np.int64)
 for i in range(max_new_tokens):
   logits, *present_key_values = decoder_session.run(None, dict(
       input_ids=input_ids,
-      attention_mask=attention_mask,
       position_ids=position_ids,
       **past_key_values,
   ))
 
   ## Update values for next generation loop
   input_ids = logits[:, -1].argmax(-1, keepdims=True)
-  attention_mask = np.ones_like(input_ids)
   position_ids = position_ids[:, -1:] + 1
   for j, key in enumerate(past_key_values):
     past_key_values[key] = present_key_values[j]
@@ -145,5 +146,4 @@ const messages = [
 // Generate a response
 const output = await generator(messages, { max_new_tokens: 512, do_sample: false });
 console.log(output[0].generated_text.at(-1).content);
-```
-
+```