Update README.md

README.md

@@ -16,13 +16,33 @@ This model is a fine-tuned version of [bigscience/bloom-560m](https://huggingfac
 ## Usage
 
 ```py
-from transformers import pipeline
+from datasets import load_dataset
+from transformers import BloomTokenizerFast, BloomForCausalLM
 
-model_ckpt = "mrm8488/electricidad-base-finetuned-sst2-es"
+valid_dataset = load_dataset('totto', split='validation')
 
-classifier = pipeline("sentiment-analysis", model=model_ckpt)
+from preprocess import preprocess # This file is included in the repo
 
-classifier("Here your text in Spanish!")
+# Now we linearize the tables
+valid_dataset = valid_dataset.map(preprocess) 
+
+model_ckpt = "mrm8488/bloom-560m-finetuned-totto-table-to-text"
+
+tokenizer = BloomTokenizerFast.from_pretrained(ckpt)
+model = BloomForCausalLM.from_pretrained(ckpt).to("cuda")
+
+
+def explain_hl_cells(text):
+    inputs = tokenizer(text, return_tensors='pt')
+    input_ids = inputs.input_ids.to("cuda")
+    attention_mask = inputs.attention_mask.to("cuda")
+    output = model.generate(input_ids, attention_mask=attention_mask, max_length=2048, eos_token_id=tokenizer.eos_token_id) # num_beams=3, temperature=1.9
+
+    return tokenizer.decode(output[0], skip_special_tokens=False)
+
+example = valid_dataset[1]
+
+print(explain_hl_cells(example['linearized_table'])
 ``` 
 
 ### Evaluation results