Update app.py

app.py

@@ -26,7 +26,7 @@ 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: {torch.device('cuda' if torch.cuda.is_available() else 'cpu')}")
+    logger.info(f"Device: CPU")
 
 print("Starting application...")
 log_system_info()
@@ -35,53 +35,55 @@ try:
     print("Loading model and tokenizer...")
     model_id = "htigenai/finetune_test_2_4bit"
     base_model_id = "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"  # Original base model
-    
+
     with timer("Loading tokenizer"):
         try:
-            tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+            tokenizer = AutoTokenizer.from_pretrained(base_model_id, use_fast=False)
             tokenizer.pad_token = tokenizer.eos_token
             tokenizer.padding_side = "right"
         except Exception as e:
             logger.error(f"Error loading tokenizer: {str(e)}")
             raise
     logger.info("Tokenizer loaded successfully")
-    
-    # Configure quantization - using 4-bit since the base model was 4-bit
+
+    # Adjust quantization config for CPU
     bnb_config = BitsAndBytesConfig(
         load_in_4bit=True,
         bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_dtype=torch.float16,
+        bnb_4bit_compute_dtype=torch.bfloat16,  # Use bfloat16 for better CPU support
         bnb_4bit_use_double_quant=True,
     )
-    
+
     with timer("Loading model"):
         model = AutoModelForCausalLM.from_pretrained(
             model_id,
             quantization_config=bnb_config,
-            device_map="auto",
+            device_map={"": "cpu"},  # Explicitly set to CPU
             trust_remote_code=True,
         )
         model.eval()
     logger.info("Model loaded successfully")
 
-    def generate_text(prompt, max_tokens=200, temperature=0.7):
+    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, 
+                    prompt,
                     return_tensors="pt",
                     padding=True,
                     truncation=True,
                     max_length=256
-                ).to(model.device)
-            
+                )
+                inputs = inputs.to("cpu")  # Ensure inputs are on CPU
+
             with timer("Generation"):
-                with torch.inference_mode():
+                with torch.no_grad():
                     outputs = model.generate(
-                        **inputs,
+                        input_ids=inputs["input_ids"],
+                        attention_mask=inputs["attention_mask"],
                         max_new_tokens=max_tokens,
                         temperature=temperature,
                         top_p=0.95,
@@ -90,20 +92,18 @@ try:
                         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()
-                if torch.cuda.is_available():
-                    torch.cuda.empty_cache()
-            
+
             return generated_text
-            
+
         except Exception as e:
             logger.error(f"Error during generation: {str(e)}")
             return f"Error during generation: {str(e)}"
@@ -113,14 +113,14 @@ try:
         fn=generate_text,
         inputs=[
             gr.Textbox(
-                lines=3, 
+                lines=3,
                 placeholder="Enter your prompt here...",
                 label="Input Prompt"
             ),
             gr.Slider(
-                minimum=50,
-                maximum=200,
-                value=100,
+                minimum=20,
+                maximum=100,
+                value=50,
                 step=10,
                 label="Max Tokens"
             ),
@@ -134,20 +134,20 @@ try:
         ],
         outputs=gr.Textbox(
             label="Generated Response",
-            lines=5
+            lines=10
         ),
         title="HTIGENAI Reflection Analyzer - Test",
         description="Enter a prompt to generate text. This model is fine-tuned from Llama 3.1 8B Instruct.",
         examples=[
-            ["What are your thoughts about cats?", 100, 0.7],
-            ["Write a short story about a magical forest", 150, 0.8],
-            ["Explain quantum computing to a 5-year-old", 75, 0.5],
+            ["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],
         ]
     )
 
     # Launch the interface
-    iface.launch(server_name="0.0.0.0")
+    iface.launch()
 
 except Exception as e:
     logger.error(f"Application startup failed: {str(e)}")
-    raise
+    raise