token probability feature

Author: Paula Navarrete Díaz

notebooks/word_count_prototype.ipynb

@@ -101,6 +101,19 @@ def get_token_counts_as_dataframe(self):
         )
         dataframe = dataframe.sort_values("token")
         return dataframe
+    
+    def get_token_probabilities_as_dataframe(self):
+        """Returns the token probabilities of the corpus as a Pandas DataFrame with columns 'token', 'probability'
+        """
+        df=self.get_token_counts_as_dataframe()
+        total_tokens=df["count"].sum()
+            
+        dataframe = pd.DataFrame.from_records(
+            list(self.token_counter.items()), columns=["token", "probability"]
+        )
+        dataframe = dataframe.sort_values("token")
+        dataframe["probability"]=dataframe["probability"].div(total_tokens)
+        return dataframe
 
     def save_token_counts(self, csv_file):
         """Saves the counts of tokens the corpus to a specified

src/cdstemplate/word_count.py

@@ -103,3 +103,9 @@ def test_corpus_counter_save_csv(tmp_path):
     assert my_csv.is_file()
     expected_csv = "token,count\na,2\nb,1\nc,1\nx,1\ny,1\nz,1\n"
     assert my_csv.read_text() == expected_csv
+
+def test_token_probabilities_to_dataframe():
+    cc = word_count.CorpusCounter()
+    cc.add_doc("A a B b b b b b")
+    dataframe = cc.get_token_probabilities_as_dataframe()
+    assert dataframe["probability"].sum()==1