Update README.md

README.md

@@ -155,6 +155,52 @@ Users (both direct and downstream) should be made aware of the risks, biases and
 Use the code below to get started with the model.
 
 ```python
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+tokenizer = AutoTokenizer.from_pretrained('liujch1998/crystal-11b')
+model = AutoModelForSeq2SeqLM.from_pretrained('liujch1998/crystal-11b')
+model.eval()
+
+max_question_len, max_knowledge_len, max_answer_len = 128, 32, 2
+k = 1 # number of knowledge statements to generate
+top_p = 0.0001
+
+question = 'If the mass of an object gets bigger what will happen to the amount of matter contained within it? \\n (A) gets bigger (B) gets smaller'
+choices = ['A', 'B']
+
+choices_ids = tokenizer(choices, return_tensors='pt', padding='max_length', truncation='longest_first', max_length=max_answer_len).input_ids # (C, AL)
+
+prompt = question + ' \\n Knowledge: '
+prompt_tok = tokenizer(prompt, return_tensors='pt', padding='max_length', truncation='longest_first', max_length=max_question_len) # (1, QL)
+knowledges_ids = self.model.generate(
+    input_ids=prompt_tok.input_ids,
+    attention_mask=prompt_tok.attention_mask,
+    max_length=max_knowledge_len + 1,
+    min_length=3,
+    do_sample=True,
+    num_return_sequences=k,
+    top_p=top_p,
+) # (K, KL); begins with 0 ([BOS]); ends with 1 ([EOS])
+knowledges_ids = knowledges_ids[:, 1:].contiguous() # no beginning; ends with 1 ([EOS])
+knowledges = tokenizer.batch_decode(knowledges_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+
+prompts = [question + (f' \\n Knowledge: {knowledge} \\n Answer: ' if knowledge != '' else ' \\n Answer:') for knowledge in knowledges]
+prompts_tok = self.tokenizer(prompts, return_tensors='pt', padding='max_length', truncation='longest_first', max_length=max_question_len + max_knowledge_len) # (K, QL+KL)
+output = model(
+    input_ids=prompts_tok.input_ids,
+    attention_mask=prompts_tok.attention_mask,
+    labels=choices_ids[0].unsqueeze(0).repeat(len(knowledges), 1),
+)
+logitsss = output.logits # (K, AL, V)
+logitss = logitsss[:, 0, :] # (K, V)
+choice_ids = choices_ids[:, 0] # (C)
+answer_logitss = logitss.gather(dim=1, index=choice_ids.unsqueeze(0).expand(len(knowledges), -1)) # (K, C)
+answer_probss = answer_logitss.softmax(dim=1) # (K, C)
+answer_probs = answer_probss.max(dim=0).values # (C)
+pred = answer_probs.argmax(dim=0).item()
+pred = choices[pred]
+
+print(f'Question: {question}\nKnowledge: {knowledges[0]}\nAnswer: {pred}')
 ```
 
 You may also refer to <https://huggingface.co/spaces/liujch1998/crystal/blob/main/app.py#L10-L86> for implementation.