Update app.py

app.py

@@ -1,31 +1,34 @@
-import os
+import threading
+import time
 from datasets import load_dataset
-import pandas as pd
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# Step 1: Load dataset and save to CSV (auto-training data)
-def load_and_save_dataset():
-    print("Loading dataset from Hugging Face...")
-    dataset = load_dataset("HuggingFaceFW/fineweb", split="train")
-    print("Saving dataset to data.csv...")
-    dataset.to_csv("data.csv")
-    print("Done! Data saved to data.csv.")
-    return "Dataset loaded and saved to data.csv."
+# Global variable to store dataset loading status
+dataset_loaded = False
+dataset_info = "Dataset not loaded yet."
 
-# Run on startup
-load_and_save_dataset()
+def load_dataset_in_background():
+    global dataset_loaded, dataset_info
+    try:
+        dataset_info = "Loading dataset..."
+        dataset = load_dataset("HuggingFaceFW/fineweb", split="train")
+        # Save to CSV if needed
+        dataset.to_csv("data.csv")
+        dataset_info = "Dataset loaded successfully!"
+        dataset_loaded = True
+    except Exception as e:
+        dataset_info = f"Error loading dataset: {e}"
 
-# Step 2: Load GPT-2 model for inference
-model_name = "gpt2"  # or "distilgpt2" for faster inference
+# Start dataset loading in background thread
+threading.Thread(target=load_dataset_in_background, daemon=True).start()
+
+# Load GPT-2 model for inference
+model_name = "gpt2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
-
-# Create generator pipeline
 generator = pipeline('text-generation', model=model, tokenizer=tokenizer, device=-1)
 
-# Function to generate responses
 def generate_response(prompt):
     responses = generator(
         prompt,
@@ -38,24 +41,24 @@ def generate_response(prompt):
     )
     return responses[0]['generated_text'].strip()
 
-# Step 3: Build Gradio interface
+# Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("## GPT-2 Based AI Assistant with Dataset Loaded from Hugging Face")
-    gr.Textbox(value="Loading dataset...", interactive=False, lines=2)
-    fetch_button = gr.Button("Load Dataset and Save CSV")
-    output_message = gr.Textbox()
+    gr.Markdown("## AI Assistant with Background Dataset Loading")
+    dataset_status = gr.Textbox(value=dataset_info, label="Dataset Loading Status", interactive=False, lines=2)
 
-    def fetch_and_confirm():
-        msg = load_and_save_dataset()
-        return msg
+    def get_dataset_status():
+        return dataset_info
 
-    fetch_button.click(fetch_and_confirm, outputs=output_message)
+    # Refresh status button (or auto-update)
+    refresh_btn = gr.Button("Check Dataset Status")
+    refresh_btn.click(get_dataset_status, outputs=dataset_status)
 
-    gr.Markdown("### Ask the AI Assistant")
-    prompt_input = gr.Textbox(label="Enter your prompt", placeholder="Say something...")
-    response_output = gr.Textbox(label="Response", lines=10)
+    gr.Markdown("### Chat with the AI")
+    prompt_input = gr.Textbox(label="Your prompt", placeholder="Ask me anything...")
+    response_output = gr.Textbox(label="AI Response", lines=10)
 
     def respond(prompt):
+        # You can implement logic to use dataset info here if needed
         return generate_response(prompt)
 
     gr.Button("Ask").click(respond, inputs=prompt_input, outputs=response_output)