Update app.py

app.py

@@ -3,34 +3,51 @@ import spaces  # ZeroGPU helper module
 from transformers import pipeline
 
 # Preload the text-generation model on CPU at startup.
-# Model: EleutherAI/gpt-j-6B (https://huggingface.co/EleutherAI/gpt-j-6B)
-# We load on CPU (device=-1) so that initialization is done before the GUI is up.
+# We load EleutherAI/gpt-j-6B on CPU (device=-1).
 generator = pipeline("text-generation", model="EleutherAI/gpt-j-6B", device=-1)
 
-@spaces.GPU  # This decorator ensures ZeroGPU allocates a GPU only during this function call.
+@spaces.GPU  # ZeroGPU will allocate the GPU only during this function call.
 def expand_prompt(prompt, num_variants=5, max_length=100):
     """
-    Given a basic prompt, generate `num_variants` expanded prompt variants.
-    Before generation, the model is moved to GPU (A100), and after generation it's moved back to CPU.
+    Given a basic prompt, generate `num_variants` expanded prompts using GPT-J-6B.
+    This function explicitly tokenizes the input with truncation (strategy 'longest_first'),
+    moves the input to GPU, generates output using the GPU, and then moves the model back to CPU.
     """
-    # Move the model to GPU for generation.
+    # Move model to GPU for generation.
     generator.model.to("cuda")
-    outputs = generator(prompt, max_length=max_length, num_return_sequences=num_variants, do_sample=True)
-    # Move the model back to CPU after generation.
+    
+    # Explicitly tokenize the input with truncation.
+    inputs = generator.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=max_length)
+    # Move inputs to GPU.
+    inputs = {k: v.to("cuda") for k, v in inputs.items()}
+    
+    # Generate text, explicitly setting pad_token_id to eos_token_id.
+    outputs = generator.model.generate(
+        **inputs,
+        max_length=max_length,
+        num_return_sequences=num_variants,
+        do_sample=True,
+        pad_token_id=generator.tokenizer.eos_token_id
+    )
+    
+    # Decode outputs.
+    expanded = [generator.tokenizer.decode(output, skip_special_tokens=True).strip() for output in outputs]
+    
+    # Move model back to CPU.
     generator.model.to("cpu")
-    expanded = [out["generated_text"].strip() for out in outputs]
+    
     return "\n\n".join(expanded)
 
-# Create a Gradio Interface
 iface = gr.Interface(
     fn=expand_prompt,
     inputs=gr.Textbox(lines=2, placeholder="Enter your basic prompt here...", label="Basic Prompt"),
     outputs=gr.Textbox(lines=10, label="Expanded Prompts"),
     title="Prompt Expansion Generator",
     description=(
-        "Enter a basic prompt and receive 5 creative, expanded prompt variants. "
-        "The model (EleutherAI/gpt-j-6B) is preloaded on CPU at startup and then moved to GPU (via ZeroGPU) only "
-        "when a prompt is submitted. Simply copy the output for use with your downstream image-generation pipeline."
+        "Enter a basic prompt to receive 5 creative, expanded prompt variants. "
+        "The model (EleutherAI/gpt-j-6B) is preloaded on CPU at startup and moved to GPU (via ZeroGPU) for generation. "
+        "Input is tokenized with truncation enabled. Once generation is complete, the model is moved back to CPU. "
+        "Simply copy the output for use in your downstream image-generation pipeline."
     )
 )