Update app.py

app.py

@@ -4,112 +4,103 @@ import torch
 import logging
 import sys
 import gc
-import time
 from contextlib import contextmanager
 
 # Set up logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(levelname)s - %(message)s',
-    handlers=[logging.StreamHandler(sys.stdout)]
-)
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-@contextmanager
-def timer(description: str):
-    start = time.time()
-    yield
-    elapsed = time.time() - start
-    logger.info(f"{description}: {elapsed:.2f} seconds")
-
-def log_system_info():
-    """Log system information for debugging"""
-    logger.info(f"Python version: {sys.version}")
-    logger.info(f"PyTorch version: {torch.__version__}")
-    logger.info(f"Device: CPU")
-
-print("Starting application...")
-log_system_info()
+logger.info("Starting application...")
+logger.info(f"CUDA available: {torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    logger.info(f"GPU: {torch.cuda.get_device_name(0)}")
 
 try:
-    print("Loading model and tokenizer...")
-    
-    model_id = "htigenai/finetune_test"  # Replace with your chosen model ID
-
-    with timer("Loading tokenizer"):
-        tokenizer = AutoTokenizer.from_pretrained(
-            model_id,
-            use_fast=True,  # Use fast tokenizer for better performance
-            cache_dir='./cache'
-        )
-        tokenizer.pad_token = tokenizer.eos_token
+    logger.info("Loading tokenizer...")
+    model_id = "htigenai/finetune_test_2"
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_id,
+        use_fast=False  # Use slow tokenizer to save memory
+    )
+    tokenizer.pad_token = tokenizer.eos_token
     logger.info("Tokenizer loaded successfully")
 
-    with timer("Loading model"):
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            device_map={"": "cpu"},
-            cache_dir='./cache'
-        )
-        model.eval()
-    logger.info("Model loaded successfully")
+    logger.info("Loading model in 8-bit...")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        device_map="auto",
+        load_in_8bit=True,  # Load in 8-bit instead of 4-bit
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True,
+        max_memory={0: "12GB", "cpu": "4GB"}  # Limit memory usage
+    )
+    model.eval()
+    logger.info("Model loaded successfully in 8-bit")
+
+    # Clear any residual memory
+    gc.collect()
+    torch.cuda.empty_cache()
 
     def generate_text(prompt, max_tokens=100, temperature=0.7):
-        """Generate text based on the input prompt."""
         try:
-            logger.info(f"Starting generation for prompt: {prompt[:50]}...")
-
-            with timer("Tokenization"):
-                inputs = tokenizer(
-                    prompt,
-                    return_tensors="pt",
-                    padding=True,
-                    truncation=True,
-                    max_length=256
-                ).to("cpu")  # Ensure inputs are on CPU
-
-            with timer("Generation"):
-                with torch.no_grad():
-                    outputs = model.generate(
-                        input_ids=inputs["input_ids"],
-                        attention_mask=inputs["attention_mask"],
-                        max_new_tokens=max_tokens,
-                        temperature=temperature,
-                        top_p=0.95,
-                        do_sample=True,
-                        pad_token_id=tokenizer.pad_token_id,
-                        eos_token_id=tokenizer.eos_token_id,
-                        repetition_penalty=1.1,
-                    )
-
-            with timer("Decoding"):
-                generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-            logger.info("Text generation completed successfully")
-
-            # Clean up
-            with timer("Cleanup"):
-                gc.collect()
-
-            return generated_text
+            # Format the prompt
+            formatted_prompt = f"### Human: {prompt}\n\n### Assistant:"
+            
+            # Generate with memory-efficient settings
+            inputs = tokenizer(
+                formatted_prompt,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=256  # Limit input length
+            ).to(model.device)
+
+            with torch.inference_mode():
+                outputs = model.generate(
+                    **inputs,
+                    max_new_tokens=max_tokens,
+                    temperature=temperature,
+                    do_sample=True,
+                    top_p=0.95,
+                    repetition_penalty=1.2,
+                    pad_token_id=tokenizer.pad_token_id,
+                    eos_token_id=tokenizer.eos_token_id,
+                    early_stopping=True,
+                    no_repeat_ngram_size=3,
+                    use_cache=True
+                )
+
+            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            # Extract only the assistant's response
+            if "### Assistant:" in response:
+                response = response.split("### Assistant:")[-1].strip()
+            
+            # Clean up memory after generation
+            del outputs, inputs
+            gc.collect()
+            torch.cuda.empty_cache()
+            
+            return response
 
         except Exception as e:
             logger.error(f"Error during generation: {str(e)}")
-            return f"Error during generation: {str(e)}"
+            return f"Error generating response: {str(e)}"
 
-    # Create Gradio interface
+    # Create a more memory-efficient Gradio interface
     iface = gr.Interface(
         fn=generate_text,
         inputs=[
             gr.Textbox(
-                lines=3,
+                lines=3, 
                 placeholder="Enter your prompt here...",
-                label="Input Prompt"
+                label="Input Prompt",
+                max_lines=5
             ),
             gr.Slider(
-                minimum=20,
-                maximum=200,
-                value=100,
+                minimum=10,
+                maximum=100,
+                value=50,
                 step=10,
                 label="Max Tokens"
             ),
@@ -123,19 +114,26 @@ try:
         ],
         outputs=gr.Textbox(
             label="Generated Response",
-            lines=10
+            lines=5
         ),
-        title="Text Generation Demo",
-        description="Enter a prompt to generate text.",
+        title="HTIGENAI Reflection Analyzer (8-bit)",
+        description="8-bit quantized text generation. Please keep prompts concise for best results.",
         examples=[
-            ["What are your thoughts about cats?", 50, 0.7],
-            ["Write a short story about a magical forest", 60, 0.8],
-            ["Explain quantum computing to a 5-year-old", 40, 0.5],
-        ]
+            ["What is machine learning?", 50, 0.7],
+            ["Explain quantum computing", 50, 0.7],
+        ],
+        cache_examples=False
     )
 
-    iface.launch()
+    # Launch with minimal memory usage
+    iface.launch(
+        server_name="0.0.0.0",
+        share=False,
+        show_error=True,
+        enable_queue=True,
+        max_threads=1
+    )
 
 except Exception as e:
     logger.error(f"Application startup failed: {str(e)}")
-    raise
+    raise