Update README.md

README.md

@@ -14,12 +14,16 @@ language:
 This is a [TRL language model](https://github.com/huggingface/trl) that has been fine-tuned with reinforcement learning to
  guide the model outputs according to a simulated human feedback. The model was fine-tuned for classification of cancer / diabetes based on clinical notes.
  
+```bash
+pip install torch transformers trl peft
+```
 
 ## Usage
 
 ```python
 from transformers import AutoTokenizer
 from trl import AutoModelForCausalLMWithValueHead
+from peft import LoraConfig
 
 tokenizer_kwargs = {
   "padding": "max_length",
@@ -42,11 +46,22 @@ generation_kwargs = {
   "repetition_penalty":1.2
 }
 
-model = AutoModelForCausalLMWithValueHead.from_pretrained("hanyinwang/layer-project-diagnostic-mistral")
+model = AutoModelForCausalLMWithValueHead.from_pretrained("hanyinwang/layer-project-diagnostic-mistral").cuda()
 
-query_tensors = tokenizer.encode(<prompt>, return_tensors="pt")
-prompt_length = input_ids.shape[1]
+def format_prompt_mistral(text, condition):
+    prompt = """<s>[INST]You are a medical doctor specialized in %s diagnosis. 
+From the provided document, assert if the patient historically and currently has %s.
+For each condition, only pick from "YES", "NO", or "MAYBE". And you must follow format without anything further. The results have to be directly parseable with python json.loads(). 
+Sample output: {"%s": "MAYBE"}
+Never output anything beyond the format.[/INST]
+Provided document: %s"""%(condition, condition, condition, text)
+    return prompt
 
-outputs = model(query_tensors, **generation_kwargs)
-response = tokenizer.decode(outputs[prompt_length:])
+query_tensors = tokenizer.encode(format_prompt_mistral(<note>, <condition>), return_tensors="pt")
+# <note>: clinical note
+# <condition>: "cancer" or "diabetes"
+prompt_length = query_tensors.shape[1]
+
+outputs = model.generate(query_tensors.cuda(), **generation_kwargs)
+response = tokenizer.decode(outputs[0][prompt_length:])
 ```