Update app.py

app.py

@@ -1,22 +1,65 @@
 import gradio as gr
 import spaces
-from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
+from datasets import load_dataset
+from transformers import (
+    AutoConfig,
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    DataCollatorForLanguageModeling,
+    Trainer,
+    TrainingArguments,
+    pipeline
+)
 
-text_pipeline = None  # global var to hold our pipeline once loaded
+#############################################################
+# ZeroGPU REQUIREMENT:
+#   - No CUDA references at global scope.
+#   - All GPU usage within @spaces.GPU(...) functions.
+#############################################################
 
-@spaces.GPU(duration=120)  # request up to 120s GPU time to load the model
-def load_model():
+# 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
+
+# We'll store the "inference pipeline" after training
+TEXT_PIPELINE = None
+
+@spaces.GPU(duration=300)  # up to 5 minutes for a mini-finetraining
+def finetune_small_subset():
     """
-    This function will run in a *child* process that has GPU allocated.
-    We can safely do device_map="auto" or .to("cuda") here.
+    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.
     """
+
+    # -------------------------------
+    # A) Load a small 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)
+    # -------------------------------
     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
     )
@@ -24,41 +67,146 @@ def load_model():
         "wuhp/myr1",
         subfolder="myr1",
         config=config,
-        torch_dtype="auto",  # triggers GPU usage
-        device_map="auto",   # triggers GPU usage
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto", 
         trust_remote_code=True
     )
-    text_pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-    return text_pipe
+
+    # -------------------------------
+    # C) Process dataset
+    # -------------------------------
+    ds = ds.map(format_and_tokenize, batched=True, remove_columns=["text"])
+    ds.set_format("torch")
+
+    # -------------------------------
+    # D) Data Collator
+    # -------------------------------
+    collator = DataCollatorForLanguageModeling(
+        tokenizer=tokenizer,
+        mlm=False
+    )
+
+    # -------------------------------
+    # E) Training Arguments + Trainer
+    # -------------------------------
+    training_args = TrainingArguments(
+        output_dir="finetuned_myr1",
+        num_train_epochs=1,  # 1 epoch for demonstration
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=2,
+        logging_steps=10,
+        save_steps=999999,        # effectively "don't save mid-training"
+        save_total_limit=1,
+        fp16=torch.cuda.is_available(),
+        # ZeroGPU ephemeral environment => no real advantage to push_to_hub
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=ds,
+        data_collator=collator,
+    )
+
+    # -------------------------------
+    # F) Train
+    # -------------------------------
+    trainer.train()
+
+    # -------------------------------
+    # G) Save local checkpoint
+    # -------------------------------
+    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.
+    finetuned_model = AutoModelForCausalLM.from_pretrained(
+        "finetuned_myr1",
+        torch_dtype=torch.float16 if torch.cuda.is_available() 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!"
 
 def ensure_pipeline():
     """
-    If we've never loaded the pipeline, call load_model() now.
-    If ZeroGPU has deallocated it, we might need to reload again.
+    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.
     """
-    global text_pipeline
-    if text_pipeline is None:
-        text_pipeline = load_model()  # <-- calls the GPU-wrapped function
-    return text_pipeline
+    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
+            trust_remote_code=True
+        )
+    return TEXT_PIPELINE
 
-@spaces.GPU(duration=60)  # up to 60s for each generate call
-def predict(prompt, max_new_tokens=64):
+@spaces.GPU(duration=120)  # up to 2 minutes to generate text
+def predict(prompt, min_new_tokens=260, max_new_tokens=2600):
     """
-    Called when the user clicks 'Generate'; ensures the model is loaded, 
-    then runs inference on GPU.
+    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.
     """
-    pipe = ensure_pipeline()
-    outputs = pipe(prompt, max_new_tokens=int(max_new_tokens))
+    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(
+        prompt,
+        min_new_tokens=int(min_new_tokens),
+        max_new_tokens=int(max_new_tokens),
+        temperature=0.7,
+        top_p=0.9
+    )
     return outputs[0]["generated_text"]
 
-# Build the Gradio UI
+#############################################################
+# Build a Gradio UI
+#############################################################
 with gr.Blocks() as demo:
-    gr.Markdown("# ZeroGPU Inference Demo")
-    prompt = gr.Textbox(label="Prompt")
-    max_tok = gr.Slider(1, 256, value=64, step=1, label="Max New Tokens")
-    output = gr.Textbox(label="Generated Text")
+    gr.Markdown("## ZeroGPU Finetuning & Long-Text Generation Demo")
 
-    generate_btn = gr.Button("Generate")
-    generate_btn.click(fn=predict, inputs=[prompt, max_tok], outputs=output)
+    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)
+
+    gr.Markdown(
+        "Once finetuning completes, or if you skip it, you can still do inference "
+        "with either the new or original model."
+    )
+
+    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.click(
+        fn=predict,
+        inputs=[prompt_in, min_tok_slider, max_tok_slider],
+        outputs=output_box
+    )
 
 demo.launch()