Update README.md

README.md

@@ -27,36 +27,39 @@ The backbone model of COSMO is the [lm-adapted T5](https://huggingface.co/google
 
 ### How to use
 
+> 💡 <b>Note:</b> The HuggingFace inference API for Cosmo is not working correctly, we gently guide you to [our repository](https://hyunw.kim/sodaverse) to try out the demo code!
+
+Below is a simple code snippet to get Cosmo running :)
+
 ```python
+import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 tokenizer = AutoTokenizer.from_pretrained("allenai/cosmo-xl")
-model = AutoModelForSeq2SeqLM.from_pretrained("allenai/cosmo-xl")
-# model.to('cuda')
+model = AutoModelForSeq2SeqLM.from_pretrained("allenai/cosmo-xl").to(device)
 
-def set_input(narrative, instruction, dialogue_history):
-    input_text = " <turn> ".join(dialogue_history)
+def set_input(situation_narrative, role_instruction, conversation_history):
+    input_text = " <turn> ".join(conversation_history)
 
-    if instruction != "":
-        input_text = instruction + " <sep> " + input_text
+    if role_instruction != "":
+        input_text = "{} <sep> {}".format(role_instruction, input_text)
 
-    if narrative != "":
-        input_text = narrative + " <sep> " + input_text
-    
+    if situation_narrative != "":
+        input_text = "{} <sep> {}".format(situation_narrative, input_text)
 
     return input_text
 
-def generate(narrative, instruction, dialogue_history):
+def generate(situation_narrative, role_instruction, conversation_history):
     """
-    narrative: the description of situation/context with the characters included (e.g., "David goes to an amusement park")
-    instruction: the perspective/speaker instruction (e.g., "Imagine you are David and speak to his friend Sarah").
-    dialogue history: the previous utterances in the dialogue in a list
+    situation_narrative: the description of situation/context with the characters included (e.g., "David goes to an amusement park")
+    role_instruction: the perspective/speaker instruction (e.g., "Imagine you are David and speak to his friend Sarah").
+    conversation_history: the previous utterances in the conversation in a list
     """
 
-    input_text = set_input(narrative, instruction, dialogue_history) 
+    input_text = set_input(role_narrative, role_instruction, conversation_history) 
 
-    inputs = tokenizer([input_text], return_tensors="pt")
-    # inputs = inputs.to('cuda')
+    inputs = tokenizer([input_text], return_tensors="pt").to(device)
     outputs = model.generate(inputs["input_ids"], max_new_tokens=128, temperature=1.0, top_p=.95, do_sample=True)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
 
@@ -65,11 +68,11 @@ def generate(narrative, instruction, dialogue_history):
 situation = "Cosmo had a really fun time participating in the EMNLP conference at Abu Dhabi."
 instruction = "You are Cosmo and you are talking to a friend." # You can also leave the instruction empty
 
-dialogue = [
+conversation = [
     "Hey, how was your trip to Abu Dhabi?"
 ]
 
-response = generate(situation, instruction, dialogue)
+response = generate(situation, instruction, conversation)
 print(response)
 ```