Update app.py

app.py

@@ -2,36 +2,37 @@ import gradio as gr
 import tiktoken
 import json
 
-# Function to load JSON/JSONL file and count tokens for each entry
+# Function to count tokens in the dataset based on the "messages" field
 def count_tokens(json_file, encoding_name):
     encoding = tiktoken.get_encoding(encoding_name)
     
     # Load the JSON or JSONL data
     with open(json_file.name, 'r') as f:
-        data = [json.loads(line) for line in f.readlines()] if json_file.name.endswith('.jsonl') else json.load(f)
-    
-    # Token counting for each item in the dataset
+        data = json.load(f) if json_file.name.endswith('.json') else [json.loads(line) for line in f.readlines()]
+
     token_counts = []
-    for item in data:
-        if isinstance(item, dict):
-            text = ' '.join([str(v) for v in item.values() if isinstance(v, str)])
-        else:
-            text = str(item)
-        
-        num_tokens = len(encoding.encode(text))
+    for entry in data:
+        conversation_token_count = 0
+        conversation_texts = []
+        if "messages" in entry:
+            for message in entry["messages"]:
+                content = message.get("content", "")
+                conversation_texts.append(content)
+                conversation_token_count += len(encoding.encode(content))
+
         token_counts.append({
-            'text': text,
-            'token_count': num_tokens
+            'conversation': ' '.join(conversation_texts),
+            'token_count': conversation_token_count
         })
     
     return token_counts
 
-# Gradio interface
+# Gradio interface function
 def token_counter(json_file, encoding_name):
     token_data = count_tokens(json_file, encoding_name)
     return token_data
 
-# Gradio UI
+# Gradio UI setup
 gr.Interface(
     fn=token_counter,
     inputs=[