Update app.py

app.py

@@ -1,154 +1,95 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline, set_seed
 import torch
 import os
 
-# 🔧 CPU Optimization Suite
-os.environ["OMP_NUM_THREADS"] = "4"
-torch.set_num_threads(4)
-torch.manual_seed(42)
-
-# 📦 Model Configuration
-MODEL_NAME = "openai-community/openai-gpt"
+# Create cache directory for model
 cache_dir = "./model_cache"
+os.makedirs(cache_dir, exist_ok=True)
 
-# 🧠 Load Tokenizer with Padding Fix
-tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_NAME,
+# Load text generation pipeline with GPU support
+generator = pipeline(
+    'text-generation',
+    model='openai-community/openai-gpt',
     cache_dir=cache_dir,
-    padding_side="left"
+    device=0 if torch.cuda.is_available() else -1
 )
 
-# ✅ Add pad_token if missing (required for batched generation)
-if tokenizer.pad_token is None:
-    # First add special token to tokenizer
-    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+def generate_text(prompt, max_length, num_return_sequences, seed):
+    """Generate text with optional seed control and error handling"""
+    if seed.strip():
+        try:
+            set_seed(int(seed))  # Handle user-provided seed
+        except ValueError:
+            pass  # Ignore invalid seeds
     
-    # Then resize model embeddings to accommodate new token
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        torch_dtype=torch.float32,
-        low_cpu_mem_usage=True,
-        cache_dir=cache_dir
+    # Generate with sampling for better diversity
+    results = generator(
+        prompt,
+        max_length=int(max_length),
+        num_return_sequences=int(num_return_sequences),
+        do_sample=True,
+        pad_token_id=generator.tokenizer.eos_token_id
     )
-    model.resize_token_embeddings(len(tokenizer))
     
-    # Finally set pad_token
-    tokenizer.pad_token = tokenizer.eos_token or '[PAD]'
-else:
-    # Load model normally if pad_token exists
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        torch_dtype=torch.float32,
-        low_cpu_mem_usage=True,
-        cache_dir=cache_dir
-    )
-
-# 🧠 Final model setup
-model = model.eval()
-
-def generate_response(prompt, max_new_tokens=128, temperature=0.7, top_p=0.9, num_sequences=1):
-    """Optimized for 18GB CPU with strict memory control"""
-    try:
-        # 🧮 Calculate max_input_length as multiple of 8
-        MAX_CONTEXT_LENGTH = 512
-        max_input_length = MAX_CONTEXT_LENGTH - (MAX_CONTEXT_LENGTH % 8)
-        
-        # 🛡️ Input Protection with strict length control
-        inputs = tokenizer(
-            prompt,
-            return_tensors="pt",
-            truncation=True,
-            max_length=max_input_length,  # Now guaranteed to be multiple of 8
-            padding="max_length",
-            pad_to_multiple_of=8,
-            add_special_tokens=False
-        )
-        
-        with torch.inference_mode():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=int(max_new_tokens),
-                temperature=float(temperature),
-                top_p=float(top_p),
-                do_sample=True,
-                num_return_sequences=int(num_sequences),
-                pad_token_id=tokenizer.convert_tokens_to_ids(tokenizer.pad_token),
-                eos_token_id=tokenizer.eos_token_id,
-                max_length=max_input_length + int(max_new_tokens)  # Total length guardrail
-            )
-        
-        return "\n\n".join([
-            f"📝 Result {i+1}:\n{tokenizer.decode(output, skip_special_tokens=True)}"
-            for i, output in enumerate(outputs)
-        ])
-    except Exception as e:
-        return f"🚨 CPU Memory Alert: {str(e)}\nTry shorter inputs or fewer results!"
+    # Format results with numbering
+    output = ""
+    for i, res in enumerate(results):
+        output += f"🔥 Result {i+1}:\n{res['generated_text']}\n\n{'-'*30}\n"
+    return output
 
-# 🎨 UI Layout with Gradio Blocks
-with gr.Blocks(theme="soft", title="GPT-1 Legacy Engine") as demo:
+# Create Gradio interface
+with gr.Blocks(theme="soft") as demo:
     gr.Markdown("""
-    # 🧠 Legacy GPT-1 Text Generator (CPU-Optimized Edition)
-    *Running the original transformer-based language model with surgical memory control*
-    
-    🔥 Features:
-    - Thread-limited execution for stable performance
-    - Input length protection (512 tokens)
-    - Batch generation support
-    - Temperature-controlled creativity
-    - Interactive examples
+    # GPT-1 Text Generator 🧠  
+    ⚠️ **Warning**: This model may generate harmful stereotypes or offensive content. Use responsibly.
+    Generate text using the original OpenAI GPT model (2018).
     """)
     
     with gr.Row():
-        with gr.Column(scale=2):
-            prompt = gr.Textbox(
-                label="Input Prompt",
-                placeholder="Enter your prompt here...",
-                lines=5
-            )
-            
-            with gr.Accordion("⚙️ Generation Parameters", open=False):
-                max_new_tokens = gr.Slider(
-                    minimum=32, maximum=256, value=128, step=16,
-                    label="Max New Tokens (Output Length)"
-                )
-                temperature = gr.Slider(
-                    minimum=0.1, maximum=1.5, value=0.7, step=0.1,
-                    label="Creativity Level (Temperature)"
-                )
-                top_p = gr.Slider(
-                    minimum=0.1, maximum=1.0, value=0.9, step=0.1,
-                    label="Top-p Sampling (Nucleus Filtering)"
-                )
-                num_sequences = gr.Slider(
-                    minimum=1, maximum=3, value=1, step=1,
-                    label="Number of Results to Generate"
-                )
-            
-            submit = gr.Button("✨ Generate Text", variant="primary")
-        
-        with gr.Column(scale=2):
-            output = gr.Textbox(
-                label="Generated Text",
-                lines=15,
-                show_copy_button=True
-            )
+        prompt = gr.Textbox(
+            label="Enter your prompt",
+            placeholder="Once upon a time...",
+            lines=3
+        )
+    
+    with gr.Row():
+        max_length = gr.Slider(
+            minimum=10, 
+            maximum=200, 
+            value=50, 
+            step=5,
+            label="Max Length"
+        )
+        num_return_sequences = gr.Slider(
+            minimum=1,
+            maximum=5,
+            value=3,
+            step=1,
+            label="Number of Results"
+        )
+        seed = gr.Textbox(
+            value="42",
+            label="Seed (leave blank for random)",
+            max_lines=1
+        )
     
-    submit.click(
-        fn=generate_response,
-        inputs=[prompt, max_new_tokens, temperature, top_p, num_sequences],
+    generate_btn = gr.Button("Generate Text", variant="primary")
+    output = gr.Textbox(label="Generated Text", lines=15)
+    
+    generate_btn.click(
+        fn=generate_text,
+        inputs=[prompt, max_length, num_return_sequences, seed],
         outputs=output
     )
     
-    gr.Examples(
-        examples=[
-            ["The future of artificial intelligence will"],
-            ["Explain quantum physics like I'm five"],
-            ["Write a haiku about machine learning"]
-        ],
-        inputs=prompt,
-        label="🚀 Try These Prompts"
-    )
+    gr.Markdown("""
+    ### Tips for Better Results  
+    - Use clear, specific prompts  
+    - Increase max length for longer stories  
+    - Try different seeds for varied outputs  
+    - GPT-1 has significant limitations compared to modern models
+    """)
 
-if __name__ == "__main__":
-    demo.launch()
+# Launch the app
+demo.launch()