Update README.md

README.md

@@ -1,7 +1,11 @@
+How to use
+
+```
 from FiDT5 import FiDT5
 from transformers import T5Tokenizer
 model = FiDT5.from_pretrained('Soyoung97/ListT5-base')
-texts = ["Query: When did Thomas Edison invent the light bulb?, Index: 1, Context: Lightning strike at Seoul National University", 
+texts = [
+"Query: When did Thomas Edison invent the light bulb?, Index: 1, Context: Lightning strike at Seoul National University", 
 "Query: When did Thomas Edison invent the light bulb?, Index: 2, Context: Thomas Edison tried to invent a device for car but failed",
 "Query: When did Thomas Edison invent the light bulb?, Index: 3, Context: Coffee is good for diet",
 "Query: When did Thomas Edison invent the light bulb?, Index: 4, Context: KEPCO fixes light problems",
@@ -9,7 +13,8 @@ texts = ["Query: When did Thomas Edison invent the light bulb?, Index: 1, Contex
 tok = T5Tokenizer.from_pretrained('t5-base')
 raw = tok(texts, return_tensors='pt', padding='max_length', max_length=128, truncation=True)
 input_tensors = {'input_ids': raw['input_ids'].unsqueeze(0), 'attention_mask': raw['attention_mask'].unsqueeze(0)}
-output = model.generate(**input_tensors, max_length=128, return_dict_in_generate=True, output_scores=True)
+output = model.generate(**input_tensors, max_length=7, return_dict_in_generate=True, output_scores=True)
 output_text = tok.batch_decode(output.sequences, skip_special_tokens=True) 
 output_text
->>> [3 1 4 2 5]
+>>> [3 1 4 2 5]
+```