Update app.py

app.py

@@ -1,34 +1,38 @@
 import threading
-import time
 from datasets import load_dataset
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# Global variable to store dataset loading status
+# Global variables for dataset status
 dataset_loaded = False
-dataset_info = "Dataset not loaded yet."
+dataset_status_message = "Dataset is still loading..."
+dataset_lock = threading.Lock()
 
 def load_dataset_in_background():
-    global dataset_loaded, dataset_info
+    global dataset_loaded, dataset_status_message
     try:
-        dataset_info = "Loading dataset..."
+        # Load dataset from Hugging Face
         dataset = load_dataset("HuggingFaceFW/fineweb", split="train")
-        # Save to CSV if needed
+        # Save to CSV for later use
         dataset.to_csv("data.csv")
-        dataset_info = "Dataset loaded successfully!"
-        dataset_loaded = True
+        with dataset_lock:
+            dataset_loaded = True
+            dataset_status_message = "Dataset loaded successfully!"
     except Exception as e:
-        dataset_info = f"Error loading dataset: {e}"
+        with dataset_lock:
+            dataset_loaded = False
+            dataset_status_message = f"Error loading dataset: {e}"
 
 # Start dataset loading in background thread
 threading.Thread(target=load_dataset_in_background, daemon=True).start()
 
-# Load GPT-2 model for inference
+# Load GPT-2 for inference
 model_name = "gpt2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 generator = pipeline('text-generation', model=model, tokenizer=tokenizer, device=-1)
 
+# Function to generate response
 def generate_response(prompt):
     responses = generator(
         prompt,
@@ -41,26 +45,25 @@ def generate_response(prompt):
     )
     return responses[0]['generated_text'].strip()
 
-# Gradio Interface
+# Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("## AI Assistant with Background Dataset Loading")
-    dataset_status = gr.Textbox(value=dataset_info, label="Dataset Loading Status", interactive=False, lines=2)
+    gr.Markdown("## GPT-2 AI Assistant with Background Dataset Loading")
+    status_box = gr.Textbox(value=dataset_status_message, label="Dataset Loading Status", interactive=False, lines=2)
 
-    def get_dataset_status():
-        return dataset_info
+    def refresh_status():
+        with dataset_lock:
+            return dataset_status_message
 
-    # Refresh status button (or auto-update)
-    refresh_btn = gr.Button("Check Dataset Status")
-    refresh_btn.click(get_dataset_status, outputs=dataset_status)
+    refresh_button = gr.Button("Check Dataset Status")
+    refresh_button.click(refresh_status, outputs=status_box)
 
     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
+    def chat(prompt):
         return generate_response(prompt)
 
-    gr.Button("Ask").click(respond, inputs=prompt_input, outputs=response_output)
+    gr.Button("Ask").click(chat, inputs=prompt_input, outputs=response_output)
 
 demo.launch()