diff --git a/u4/word_count.py b/u4/word_count.py
index 9451db9..b04f26d 100644
--- a/u4/word_count.py
+++ b/u4/word_count.py
@@ -1,15 +1,14 @@
 import re
 
-text = input()
-words = re.findall(r"[^\W_]+(?:'[^\W_]+)?", text.lower(), re.UNICODE)
-counts = {}
 
-for word in words:
-    if word in counts:
-        counts[word] += 1
-    else:
-        counts[word] = 1
+def count_words(text):
+    words = re.findall(r"[^\W_]+(?:'[^\W_]+)?", text.lower(), re.UNICODE)
+    result = {}
 
-for word, count in sorted(counts.items()):
-    print(f"{word}: {count}")
-    
\ No newline at end of file
+    for word in words:
+        if word in result:
+            result[word] += 1
+        else:
+            result[word] = 1
+
+    return result
\ No newline at end of file