Update README.md

README.md

@@ -22,12 +22,13 @@ Replace usernames and links for placeholders: "@user" and "http".
 If you're interested in retaining verified users which were also retained during training, you may keep the users listed [here](https://github.com/cardiffnlp/timelms/tree/main/data).
 ```python
 def preprocess(text):
-    new_text = []
-    for t in text.split(" "):
-        t = '@user' if t.startswith('@') and len(t) > 1 else t
-        t = 'http' if t.startswith('http') else t
-        new_text.append(t)
-    return " ".join(new_text)
+    preprocessed_text = []
+    for t in text.split():
+        if len(t) > 1:
+            t = '@user' if t[0] == '@' and t.count('@') == 1 else t
+            t = 'http' if t.startswith('http') else t
+        preprocessed_text.append(t)
+    return ' '.join(preprocessed_text)
 ```
 
 ## Example Masked Language Model 
@@ -39,8 +40,8 @@ MODEL = "cardiffnlp/twitter-roberta-base-sep2021"
 fill_mask = pipeline("fill-mask", model=MODEL, tokenizer=MODEL)
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
 
-def print_candidates():
-    for i in range(5):
+def pprint(candidates, n):
+    for i in range(n):
         token = tokenizer.decode(candidates[i]['token'])
         score = candidates[i]['score']
         print("%d) %.5f %s" % (i+1, score, token))
@@ -50,11 +51,12 @@ texts = [
     "I keep forgetting to bring a <mask>.",
     "Looking forward to watching <mask> Game tonight!",
 ]
+
 for text in texts:
     t = preprocess(text)
     print(f"{'-'*30}\n{t}")
     candidates = fill_mask(t)
-    print_candidates()
+    pprint(candidates, 5)
 ```
 
 Output: 
@@ -90,13 +92,12 @@ import numpy as np
 from scipy.spatial.distance import cosine
 from collections import Counter
 
-def get_embedding(text):
+def get_embedding(text):  # naive approach for demonstration
   text = preprocess(text)
   encoded_input = tokenizer(text, return_tensors='pt')
   features = model(**encoded_input)
   features = features[0].detach().cpu().numpy() 
-  features_mean = np.mean(features[0], axis=0) 
-  return features_mean
+  return np.mean(features[0], axis=0) 
 
 
 MODEL = "cardiffnlp/twitter-roberta-base-sep2021"