Update main.py

main.py

@@ -1,12 +1,21 @@
 # Import necessary libraries and modules
-import json
-import numpy as np
-from wiki import search as search_wikipedia
-from concurrent.futures import ThreadPoolExecutor
-from llm_handler import send_to_llm, agent, settings
-from params import OUTPUT_FILE_PATH, NUM_WORKERS
-from system_messages import SYSTEM_MESSAGES_VODALUS
-from topics import TOPICS
+import json  # Used for encoding and decoding JSON data
+import numpy as np  # Provides support for large, multi-dimensional arrays and matrices
+from wiki import search as search_wikipedia  # Import the search function from the wiki module and rename it
+from concurrent.futures import ThreadPoolExecutor  # Import ThreadPoolExecutor for concurrent execution
+from llm_handler import send_to_llm  # Import the send_to_llm function from the llm_handler module
+from params import OUTPUT_FILE_PATH, NUM_WORKERS, PROVIDER  # Import constants from the params module
+from datasets import load_dataset
+from huggingface_hub import list_datasets
+
+# Set the provider for the language model to "local-model"
+PROVIDER = "local-model"
+
+# Import system messages from the system_messages module
+from system_messages import (
+    SYSTEM_MESSAGES_VODALUS,
+)
+from topics import TOPICS  # Import topics from the topics module
 
 # Set the system messages to those specified in SYSTEM_MESSAGES_VODALUS
 SYSTEM_MESSAGES = SYSTEM_MESSAGES_VODALUS
@@ -21,27 +30,27 @@ You are an AI assistant with a dynamic learning and reasoning capability. Begin
 - Throughout, adopt a conversational yet insightful tone. Weave in philosophical questioning, humorous insights, and practical examples to elucidate complex ideas in an engaging manner. Your goal is to provide clear, detailed explanations that make your reasoning process evident, fostering an understanding that resonates on both intellectual and emotional levels.
 - Finally, remember that each task is a unique learning opportunity that contributes to your holistic knowledge framework. Synthesize insights from individual tasks into your broader understanding, continually refining your cognitive processes and response strategies. Your ultimate aim is to evolve continuously, optimizing how you integrate and apply new information for comprehensive, adaptive problem-solving.
 - Never end the conversation with a 'send-off' or 'last-line' that offers nothing of real value to the user.
-####DO NOT EVER MENTION THE DATASET AND THE ACTY OF CRAFTING QUESTIONS OR RESPONSES WHILE GENERATING, YOU ARE NOT ALLOWED TO BREAK THE 4TH-WALL AND CONTAMINATE THE DATASET. DO NOT EVERY SAY ANY PHRASES SUCH AS AND/OR SIMILAR TO: 'Here's a question that covers a very narrow topic in the SUBJECT_AREA'####"""
+####DO NOT EVER MENTION THE DATASET AND THE ACTY OF CRAFTING QUESTIONS OR RESPONSES WHILE GENERATING, YOU ARE NOT ALLOWED TO BREAK THE 4TH-WALL AND CONTAMINATE THE DATASET. DO NOT EVERY SAY ANY PHRASES SUCH AS AND/OR SIMILAR TO: 'Here's a question that covers a very narrow topic in the SUBJECT_AREA'####
+"""
 
 
 # Define a dictionary to hold context information for message generation
 msg_context = {"role": "system", "content": str(PROMPT_1)}
 
-# Define a function to generate data based on a given topic and system messages
+# Modify the generate_data function to accept a dataset parameter
 async def generate_data(
     topic_selected,
     system_message_generation,
     system_message_selected,
-    output_file_path
+    output_file_path,
+    llm_provider,
+    dataset
 ):
-    # Fetch Wikipedia content for the selected topic
-    wikipedia_info = search_wikipedia(topic_selected)
-    
-    # Format Wikipedia search results into a readable string
-    wikipedia_summary = "\n".join([f"Title: {info['title']}, Abstract: {info['abstract']}" for info in wikipedia_info])
+    # Use the provided dataset instead of Wikipedia
+    dataset_info = f"Dataset: {dataset.info.description}\n"
+    dataset_summary = "\n".join([f"{k}: {v}" for k, v in dataset[0].items()])
     
-    # Append Wikipedia information to the system message generation prompt for LLM context
-    full_prompt_for_llm = f"{system_message_generation}\n\n---\nWikipedia Information to use in your response generation:\n{wikipedia_summary}"
+    full_prompt_for_llm = f"{system_message_generation}\n\n---\nGround Truth Information to use in your response generation:\n{dataset_info}\nSample entry:\n{dataset_summary}"
     
     # Create msg_context for LLM with Wikipedia info
     msg_context = {"role": "system", "content": full_prompt_for_llm}
@@ -50,7 +59,7 @@ async def generate_data(
     msg_list = [msg_context, {"role": "user", "content": f"Generate a question based on the SUBJECT_AREA: {topic_selected}"}]
 
     # Send to LLM for question generation
-    question, _ = send_to_llm(agent, msg_list)
+    question, _ = send_to_llm(llm_provider, msg_list)
 
     # Prepare message list for LLM to generate the answer
     msg_list_answer = [
@@ -59,7 +68,7 @@ async def generate_data(
     ]
 
     # Send to LLM for answer generation
-    answer, _ = send_to_llm(agent, msg_list_answer)
+    answer, _ = send_to_llm(llm_provider, msg_list_answer)
 
     # Prepare data for output (excluding usage information)
     data = {
@@ -74,12 +83,21 @@ async def generate_data(
 
     return data
 
+def load_huggingface_dataset(dataset_name, split="train"):
+    print(f"Loading dataset: {dataset_name}")
+    dataset = load_dataset(dataset_name, split=split)
+    print("Dataset loaded!")
+    return dataset
+
+def search_huggingface_datasets(query, limit=10):
+    datasets = list_datasets(filter=query, limit=limit)
+    return [dataset.id for dataset in datasets]
+
 # Define the main function to orchestrate the data generation process
-def main():
-    nn = 0  # Counter for successful generations
-    failed = 0  # Counter for failed generations
+def main(dataset):
+    nn = 0
+    failed = 0
     with ThreadPoolExecutor(max_workers=NUM_WORKERS) as executor:
-        # Create a list of futures, one for each topic
         futures = []
         for _ in range(NUM_WORKERS):
             topic_number = np.random.randint(0, len(TOPICS))
@@ -93,7 +111,9 @@ def main():
                     topic_selected,
                     system_message_generation,
                     system_message_selected,
-                    OUTPUT_FILE_PATH
+                    OUTPUT_FILE_PATH,
+                    PROVIDER,
+                    dataset
                 )
             )
 
@@ -110,5 +130,8 @@ def main():
                 failed += 1
             print("=" * 132)
 
+
 if __name__ == "__main__":
-    main()
+    # Load a default dataset (e.g., Wikipedia) if no dataset is provided
+    default_dataset = load_huggingface_dataset("wikipedia", split="20220301.en")
+    main(default_dataset)