Update app.py

app.py

@@ -12,200 +12,172 @@ from transformers import (
     pipeline
 )
 
-#############################################################
-# ZeroGPU REQUIREMENT:
-#   - No CUDA references at global scope.
-#   - All GPU usage within @spaces.GPU(...) functions.
-#############################################################
+##############################################################################
+# ZeroGPU constraints:
+#  1) No GPU calls in top-level code
+#  2) Decorate GPU-using functions with @spaces.GPU(...)
+##############################################################################
 
-# We'll do a small subset of WikiText-2 for demonstration.
-# Real finetuning on the entire dataset likely exceeds typical ZeroGPU time.
-NUM_EXAMPLES = 1000  # or fewer to keep it quick
+TEXT_PIPELINE = None  # We'll store an inference pipeline after training (if any).
 
-# We'll store the "inference pipeline" after training
-TEXT_PIPELINE = None
+# We'll train on a subset of WikiText-2 to keep it short for ZeroGPU demonstration.
+NUM_EXAMPLES = 1000
 
-@spaces.GPU(duration=300)  # up to 5 minutes for a mini-finetraining
+@spaces.GPU(duration=300)  # 5 minutes to do a quick demo train
 def finetune_small_subset():
     """
-    1) Loads the model & tokenizer from 'wuhp/myr1'.
-    2) Loads a small subset of WikiText-2 for language modeling.
-    3) Runs a quick 1-epoch finetune.
-    4) Saves model + tokenizer to 'finetuned_myr1'.
-    5) Loads the newly trained model back into a text-generation pipeline.
-    Returns a success message.
+    Demonstration: 
+      - Loads 'wuhp/myr1' (DeepSeek) 
+      - Finetunes on a small subset of WikiText-2
+      - Disables fp16 to avoid "Attempting to unscale FP16 gradients" error
+      - Saves model to 'finetuned_myr1'
+      - Reloads as pipeline for inference
     """
-
-    # -------------------------------
-    # A) Load a small dataset
-    # -------------------------------
+    # 1) Load dataset
     ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
-    # Keep only a subset so we don't exceed time. 
     ds = ds.select(range(min(NUM_EXAMPLES, len(ds))))
 
-    def format_and_tokenize(ex):
-        # For standard LM, we just treat each line as text
-        return tokenizer(ex["text"], truncation=True, max_length=512)
-
-    # We'll define them once we have the tokenizer below.
-
-    # -------------------------------
-    # B) Load config, tokenizer, model from HF
-    #    (trust_remote_code = True for custom modeling_deepseek)
-    # -------------------------------
+    # 2) Load config, tokenizer, model
     config = AutoConfig.from_pretrained(
-        "wuhp/myr1", 
+        "wuhp/myr1",
         subfolder="myr1",
         trust_remote_code=True
     )
     tokenizer = AutoTokenizer.from_pretrained(
-        "wuhp/myr1", 
+        "wuhp/myr1",
         subfolder="myr1",
         trust_remote_code=True
     )
+
+    # If your GPU supports BF16 (e.g. A100), you can try:
+    #   bf16 = True, and fp16 = False
+    # Otherwise, just keep fp16=False
+    # We'll do bf16=False so we definitely skip half-precision
+    # (which avoids the "Attempting to unscale FP16 gradients" error).
+    bf16 = False
+    fp16 = False
+
     model = AutoModelForCausalLM.from_pretrained(
         "wuhp/myr1",
         subfolder="myr1",
         config=config,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        device_map="auto", 
+        # Only auto-detect if we do normal float32 or bfloat16. 
+        # (We do not want normal fp16 in training.)
+        torch_dtype=torch.bfloat16 if bf16 else torch.float32,
+        device_map="auto",
         trust_remote_code=True
     )
 
-    # -------------------------------
-    # C) Process dataset
-    # -------------------------------
-    ds = ds.map(format_and_tokenize, batched=True, remove_columns=["text"])
+    # 3) Tokenize data
+    def tokenize_fn(ex):
+        return tokenizer(ex["text"], truncation=True, max_length=512)
+
+    ds = ds.map(tokenize_fn, batched=True, remove_columns=["text"])
     ds.set_format("torch")
 
-    # -------------------------------
-    # D) Data Collator
-    # -------------------------------
-    collator = DataCollatorForLanguageModeling(
-        tokenizer=tokenizer,
-        mlm=False
-    )
+    collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
-    # -------------------------------
-    # E) Training Arguments + Trainer
-    # -------------------------------
+    # 4) TrainingArguments
     training_args = TrainingArguments(
         output_dir="finetuned_myr1",
-        num_train_epochs=1,  # 1 epoch for demonstration
+        num_train_epochs=1,
         per_device_train_batch_size=1,
         gradient_accumulation_steps=2,
         logging_steps=10,
-        save_steps=999999,        # effectively "don't save mid-training"
+        save_steps=999999,       # won't save mid-training
         save_total_limit=1,
-        fp16=torch.cuda.is_available(),
-        # ZeroGPU ephemeral environment => no real advantage to push_to_hub
+        # Turn off half precision explicitly
+        fp16=fp16,
+        bf16=bf16,
+        # If the above doesn't fix it, remove advanced features that auto uses 
+        # gradient scaling, or do more manual approach. 
     )
 
+    # 5) Build Trainer
     trainer = Trainer(
         model=model,
         args=training_args,
         train_dataset=ds,
-        data_collator=collator,
+        data_collator=collator
     )
 
-    # -------------------------------
-    # F) Train
-    # -------------------------------
+    # 6) Train
     trainer.train()
 
-    # -------------------------------
-    # G) Save local checkpoint
-    # -------------------------------
+    # 7) Save final
     trainer.save_model("finetuned_myr1")
     tokenizer.save_pretrained("finetuned_myr1")
 
-    # -------------------------------
-    # H) Reload the newly finetuned model as a pipeline
-    # -------------------------------
-    # (We do this so we can do inference in the same GPU session)
-    # However, if the pipeline is used *after* this function returns,
-    # we might need to re-load in a separate function call.
+    # 8) Reload the newly trained model as a pipeline
     finetuned_model = AutoModelForCausalLM.from_pretrained(
         "finetuned_myr1",
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        torch_dtype=torch.bfloat16 if bf16 else torch.float32,
         device_map="auto",
         trust_remote_code=True
     )
     global TEXT_PIPELINE
-    TEXT_PIPELINE = pipeline(
-        "text-generation",
-        model=finetuned_model,
-        tokenizer=tokenizer
-    )
-    return "Finetuning complete. Model reloaded for inference!"
+    TEXT_PIPELINE = pipeline("text-generation", model=finetuned_model, tokenizer=tokenizer)
+    return "Finetuning complete! Model reloaded for inference."
 
 def ensure_pipeline():
     """
-    If TEXT_PIPELINE is None (e.g., we didn't finetune yet),
-    let's just load the *original* model from wuhp/myr1
-    so that 'predict' can still run.
+    If we haven't finetuned yet, or if TEXT_PIPELINE is None, 
+    load the original model from 'wuhp/myr1' for inference.
     """
     global TEXT_PIPELINE
     if TEXT_PIPELINE is None:
-        # Load the original model for inference
-        TEXT_PIPELINE = pipeline(
-            "text-generation",
-            model="wuhp/myr1/myr1",    # subfolder syntax
+        tokenizer = AutoTokenizer.from_pretrained("wuhp/myr1", subfolder="myr1", trust_remote_code=True)
+        # We'll do float32 for inference if no BF16 or fp16.
+        model = AutoModelForCausalLM.from_pretrained(
+            "wuhp/myr1",
+            subfolder="myr1",
+            torch_dtype=torch.float32,
+            device_map="auto",
             trust_remote_code=True
         )
+        TEXT_PIPELINE = pipeline("text-generation", model=model, tokenizer=tokenizer)
     return TEXT_PIPELINE
 
-@spaces.GPU(duration=120)  # up to 2 minutes to generate text
-def predict(prompt, min_new_tokens=260, max_new_tokens=2600):
+@spaces.GPU(duration=120)  # up to 2 minutes for text generation
+def predict(prompt, temperature, top_p, min_new_tokens, max_new_tokens):
     """
-    Generate text from the (possibly finetuned) model. 
-    We default max_new_tokens to 2,600, but allow up to 5,000 in the UI slider.
-    We'll also ensure a minimum of 260 tokens.
+    Generates text from the (finetuned) pipeline or the original model.
+    Allows user to adjust temperature, top_p, and token range [260..5000].
     """
-    pipe = ensure_pipeline()  # load model if not already
-    # Use pipeline generate params. 
-    # The pipeline will handle do_sample by default. 
-    # We set a large max_new_tokens, but be careful about timeouts.
-    outputs = pipe(
+    pipe = ensure_pipeline()
+    out = pipe(
         prompt,
+        temperature=float(temperature),
+        top_p=float(top_p),
         min_new_tokens=int(min_new_tokens),
         max_new_tokens=int(max_new_tokens),
-        temperature=0.7,
-        top_p=0.9
+        do_sample=True
     )
-    return outputs[0]["generated_text"]
+    return out[0]["generated_text"]
 
-#############################################################
-# Build a Gradio UI
-#############################################################
+# Build Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## ZeroGPU Finetuning & Long-Text Generation Demo")
+    gr.Markdown("## ZeroGPU Mini-Finetuning (No FP16) + Long Text Generation")
 
-    finetune_btn = gr.Button("Finetune on a small WikiText-2 subset (5 min limit)")
-    finetune_status = gr.Textbox(label="Status")
-    # When user clicks, we run 'finetune_small_subset'
-    finetune_btn.click(fn=finetune_small_subset, outputs=finetune_status)
+    # 1) Button to run finetune_small_subset()
+    finetune_btn = gr.Button("Finetune WikiText-2 (Subset)")
+    status_box = gr.Textbox(label="Finetune Status")
+    finetune_btn.click(fn=finetune_small_subset, outputs=status_box)
 
-    gr.Markdown(
-        "Once finetuning completes, or if you skip it, you can still do inference "
-        "with either the new or original model."
-    )
+    gr.Markdown("Use 'Generate' to produce text from either the newly finetuned or original model.")
+
+    prompt_in = gr.Textbox(lines=3, label="Prompt")
+    temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="Temperature")
+    top_p = gr.Slider(0.0, 1.0, value=0.9, step=0.05, label="Top-p")
+    min_tokens = gr.Slider(260, 5000, value=260, step=10, label="Min New Tokens")
+    max_tokens = gr.Slider(260, 5000, value=500, step=50, label="Max New Tokens")
 
-    prompt_in = gr.Textbox(label="Prompt", lines=3)
-    min_tok_slider = gr.Slider(
-        minimum=260, maximum=5000, value=260, step=10,
-        label="Minimum New Tokens"
-    )
-    max_tok_slider = gr.Slider(
-        minimum=260, maximum=5000, value=2600, step=50,
-        label="Maximum New Tokens"
-    )
-    gen_btn = gr.Button("Generate")
     output_box = gr.Textbox(label="Generated Text", lines=12)
+    gen_btn = gr.Button("Generate")
 
     gen_btn.click(
         fn=predict,
-        inputs=[prompt_in, min_tok_slider, max_tok_slider],
+        inputs=[prompt_in, temperature, top_p, min_tokens, max_tokens],
         outputs=output_box
     )