newd

app.py

@@ -2,197 +2,174 @@
 
 import os
 import sys
-import asyncio
 import warnings
-import signal
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import gradio as gr
 
-# =================== ASYNCIO FIX ===================
-# Fix for the asyncio cleanup error
-if sys.version_info >= (3, 8) and sys.platform.startswith('linux'):
-    # This prevents the error on Linux with Python 3.8+
-    try:
-        import uvloop
-        uvloop.install()
-    except ImportError:
-        pass
+# =================== CONFIGURATION ===================
+MODEL_ID = "abdelac/Mistral_Test"
+USE_QUANTIZATION = True  # MUST be True for 16GB RAM
 
-# Suppress warnings
-warnings.filterwarnings("ignore")
-os.environ["PYTHONWARNINGS"] = "ignore"
+# =================== QUANTIZATION SETUP ===================
+if USE_QUANTIZATION:
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,                    # Critical for memory
+        bnb_4bit_quant_type="nf4",           # 4-bit quantization
+        bnb_4bit_compute_dtype=torch.float16, # Compute in float16
+        bnb_4bit_use_double_quant=True,      # Extra memory savings
+        llm_int8_enable_fp32_cpu_offload=True  # Offload to CPU if needed
+    )
+else:
+    bnb_config = None
 
 # =================== MODEL LOADING ===================
 @gr.cache_resource
 def load_model():
-    """Load the TinyLlama model"""
-    print("🚀 Loading Mistral_Test  model...")
-    MODEL_ID = "abdelac/Mistral_Test"
+    """Load Mistral model with quantization"""
+    print(f"🚀 Loading {MODEL_ID}...")
     
+    # Load tokenizer
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    
+    # Configure model loading based on quantization
+    load_kwargs = {
+        "torch_dtype": torch.float16,
+        "device_map": "auto",
+        "low_cpu_mem_usage": True,
+    }
+    
+    if USE_QUANTIZATION:
+        load_kwargs["quantization_config"] = bnb_config
+        print("✅ Using 4-bit quantization (~4GB RAM)")
+    else:
+        load_kwargs["device_map"] = "cpu"
+        print("⚠️ Using CPU only (slow but safe)")
+    
+    # Load model
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
-        torch_dtype=torch.float32,
-        device_map="cpu", 
-        low_cpu_mem_usage=True,
-        offload_folder="offload"
+        **load_kwargs
     )
     
+    # Set padding token if not present
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
     print("✅ Model loaded successfully!")
     return tokenizer, model
 
-# Load model once
-tokenizer, model = load_model()
-
-# =================== GENERATION FUNCTION ===================
-def generate_text(prompt, max_tokens=150, temperature=0.7):
-    """Generate text based on prompt"""
+# =================== MEMORY-EFFICIENT GENERATION ===================
+def generate_text(prompt, max_tokens=100, temperature=0.7):
+    """Generate text with memory constraints"""
     try:
-        # Tokenize
-        inputs = tokenizer(prompt, return_tensors="pt")
+        tokenizer, model = load_model()
+        
+        # Tokenize with truncation
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512
+        ).to(model.device)
         
-        # Generate
+        # Generate with conservative settings
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=max_tokens,
+                max_new_tokens=min(max_tokens, 150),  # Cap at 150
                 temperature=temperature,
                 do_sample=True,
-                pad_token_id=tokenizer.eos_token_id
+                pad_token_id=tokenizer.eos_token_id,
+                repetition_penalty=1.1,               # Prevent repetition
+                no_repeat_ngram_size=2,
+                early_stopping=True
             )
         
         # Decode
         result = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return result
         
+    except torch.cuda.OutOfMemoryError:
+        return "❌ Out of memory! Try reducing max tokens or using CPU mode."
     except Exception as e:
         return f"❌ Error: {str(e)}"
 
-# =================== GRADIO INTERFACE ===================
+# =================== SIMPLIFIED INTERFACE ===================
 def create_interface():
-    """Create the Gradio interface"""
+    """Create memory-aware interface"""
     with gr.Blocks(
-        title="🦙 Mistral_Test Demo",
-        theme=gr.themes.Soft(),
-        css=".gradio-container {max-width: 800px !important}"
+        title="🦅 Mistral Test Demo",
+        theme=gr.themes.Soft()
     ) as demo:
         
-        gr.Markdown("""
-        # 🦙 TinyLlama Text Generator
+        gr.Markdown(f"""
+        # 🦅 Mistral Test Demo
         
-        Generate text using the TinyLlama model (1.1B parameters)
+        **Model:** [{MODEL_ID}](https://huggingface.co/{MODEL_ID})
+        **Mode:** {'4-bit Quantized' if USE_QUANTIZATION else 'CPU'}
         
-        **Model**: [abdelac/tinyllama](https://huggingface.co/abdelac/tinyllama)
+        ⚠️ **Note:** Mistral 7B requires quantization to run in 16GB RAM
         """)
         
         with gr.Row():
-            with gr.Column(scale=2):
-                prompt = gr.Textbox(
-                    label="📝 Input Prompt",
-                    placeholder="Type your text here...",
-                    lines=5,
-                    value="Once upon a time in a magical forest,"
-                )
-                
-                with gr.Row():
-                    max_tokens = gr.Slider(
-                        50, 500, value=150,
-                        label="📏 Max Tokens",
-                        info="Maximum length of generated text"
-                    )
-                    temperature = gr.Slider(
-                        0.1, 2.0, value=0.7,
-                        label="🌡️ Temperature",
-                        info="Higher = more creative, Lower = more focused"
-                    )
-                
-                with gr.Row():
-                    generate_btn = gr.Button(
-                        "✨ Generate",
-                        variant="primary",
-                        size="lg"
-                    )
-                    clear_btn = gr.Button(
-                        "🗑️ Clear",
-                        variant="secondary"
-                    )
-            
-            with gr.Column(scale=3):
-                output = gr.Textbox(
-                    label="📄 Generated Text",
-                    lines=12,
-                    interactive=False
-                )
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="Enter your text...",
+                lines=3,
+                value="What is artificial intelligence?"
+            )
         
-        # Examples
-        gr.Examples(
-            examples=[
-                ["Write a short story about a robot learning to paint"],
-                ["Explain quantum computing in simple terms"],
-                ["Python function to calculate fibonacci sequence:"],
-                ["The benefits of renewable energy include"],
-                ["Write a poem about artificial intelligence"]
-            ],
-            inputs=prompt,
-            label="💡 Try these examples"
-        )
+        with gr.Row():
+            max_tokens = gr.Slider(
+                30, 150, value=80,  # Reduced max for memory
+                label="Max Tokens",
+                info="Higher values use more memory"
+            )
+            temperature = gr.Slider(
+                0.1, 1.0, value=0.7,
+                label="Temperature"
+            )
         
-        # Functions
-        generate_btn.click(
-            fn=generate_text,
-            inputs=[prompt, max_tokens, temperature],
-            outputs=output,
-            api_name="generate"
-        )
+        generate_btn = gr.Button("Generate", variant="primary", size="lg")
         
-        clear_btn.click(
-            fn=lambda: ("", ""),
-            inputs=[],
-            outputs=[prompt, output]
+        output = gr.Textbox(
+            label="Generated Text",
+            lines=8,
+            show_copy_button=True
         )
         
-        # Status
-        gr.Markdown("---")
+        # Memory warning
         gr.Markdown("""
-        <div style='text-align: center; color: #666; font-size: 0.9em;'>
-        ✅ Model loaded successfully | 🚀 Ready to generate text
-        </div>
+        ### 💡 Memory Optimization Tips:
+        1. **Max Tokens ≤ 100** for best results
+        2. **Temperature ~0.7** for balanced output
+        3. If OOM occurs, refresh the page
+        4. Close other tabs/applications
         """)
+        
+        # Connect button
+        generate_btn.click(
+            fn=generate_text,
+            inputs=[prompt, max_tokens, temperature],
+            outputs=output
+        )
     
     return demo
 
-# =================== MAIN ENTRY POINT ===================
-def main():
-    """Main function with proper cleanup"""
-    demo = create_interface()
-    
-    # Clean launch configuration
-    try:
-        demo.launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            share=False,
-            quiet=True,          # Reduce console output
-            debug=False,         # Disable debug mode
-            show_error=True,     # Show errors in UI
-            favicon_path=None,
-            ssl_verify=True,
-            max_file_size="2MB",
-            allowed_paths=["./"],
-            blocked_paths=[]
-        )
-    except KeyboardInterrupt:
-        print("\n👋 Shutting down gracefully...")
-        sys.exit(0)
-    except Exception as e:
-        print(f"❌ Error: {e}")
-        sys.exit(1)
-
+# =================== MAIN ===================
 if __name__ == "__main__":
-    # Set up signal handlers for clean shutdown
-    signal.signal(signal.SIGINT, lambda s, f: sys.exit(0))
-    signal.signal(signal.SIGTERM, lambda s, f: sys.exit(0))
+    # Suppress warnings
+    warnings.filterwarnings("ignore")
+    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
     
-    # Run the app
-    main()
+    # Create and launch
+    demo = create_interface()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        quiet=True,
+        debug=False,
+        show_error=True
+    )