dataroadmap commited on
Commit
3750513
1 Parent(s): 02cae1b

initial push

Browse files
Files changed (1) hide show
  1. whatsapp_chat_custom.py +49 -0
whatsapp_chat_custom.py ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # created custom class for WhatsAppChatLoader - because original langchain one isnt working
2
+
3
+ import re
4
+ from pathlib import Path
5
+ from typing import List
6
+
7
+ from langchain.docstore.document import Document
8
+ from langchain.document_loaders.base import BaseLoader
9
+
10
+
11
+ def concatenate_rows(date: str, sender: str, text: str) -> str:
12
+ """Combine message information in a readable format ready to be used."""
13
+ return f"{sender} on {date}: {text}\n\n"
14
+
15
+ # def concatenate_rows(date: str, sender: str, text: str) -> str:
16
+ # """Combine message information in a readable format ready to be used."""
17
+ # return f"{text}\n"
18
+
19
+ class WhatsAppChatLoader(BaseLoader):
20
+ """Load `WhatsApp` messages text file."""
21
+
22
+ def __init__(self, path: str):
23
+ """Initialize with path."""
24
+ self.file_path = path
25
+
26
+ def load(self) -> List[Document]:
27
+ """Load documents."""
28
+ p = Path(self.file_path)
29
+ text_content = ""
30
+
31
+ ignore_lines = ["This message was deleted", "<Media omitted>"]
32
+ #########################################################################################
33
+ # original code from langchain replaced with this code
34
+ #########################################################################################
35
+ # use https://whatstk.streamlit.app/ to get CSV
36
+ import pandas as pd
37
+ df = pd.read_csv(p)[['date', 'username', 'message']]
38
+
39
+ for i,row in df.iterrows():
40
+ date = row['date']
41
+ sender = row['username']
42
+ text = row['message']
43
+
44
+ if not any(x in text for x in ignore_lines):
45
+ text_content += concatenate_rows(date, sender, text)
46
+
47
+ metadata = {"source": str(p)}
48
+
49
+ return [Document(page_content=text_content.strip(), metadata=metadata)]