Update app.py

app.py

@@ -2,83 +2,75 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import gc
-from threading import Lock
+
+# Global model and tokenizer
+model = None
+tokenizer = None
 
 def load_model():
-    """Load and optimize the model for CPU usage"""
+    global model, tokenizer
     model_name = "Spestly/Athena-1-1.5B"
     
-    # Load model with CPU optimizations
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         low_cpu_mem_usage=True,
-        torch_dtype=torch.float32,  # Use float32 for CPU
+        torch_dtype=torch.float32,
         device_map="cpu"
     )
-    
-    # Load tokenizer
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    
-    # Enable memory efficient settings
-    model.eval()  # Set to evaluation mode
-    
-    return model, tokenizer
-
-# Global lock for thread safety
-generation_lock = Lock()
+    model.eval()
 
 def generate_response(input_text, history):
-    """Generate a response with CPU-optimized settings"""
+    global model, tokenizer
+    
+    # Load model if not loaded
+    if model is None or tokenizer is None:
+        load_model()
+        
     try:
-        with generation_lock:  # Ensure thread safety
-            instruction = "You are an AI assistant. Respond to the following instruction clearly and concisely:\n"
-            formatted_input = instruction + input_text
-            
-            # Generate with optimized settings for CPU
-            inputs = tokenizer(
-                formatted_input,
-                return_tensors="pt",
-                truncation=True,
-                max_length=256  # Reduced for CPU memory constraints
+        instruction = "You are an AI assistant. Respond to the following instruction clearly and concisely:\n"
+        formatted_input = instruction + input_text
+        
+        inputs = tokenizer(
+            formatted_input,
+            return_tensors="pt",
+            truncation=True,
+            max_length=256
+        )
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=100,
+                do_sample=True,
+                top_k=50,
+                top_p=0.9,
+                temperature=0.7,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                use_cache=True,
+                repetition_penalty=1.2,
+                num_beams=1
             )
-            
-            with torch.no_grad():  # Memory efficiency
-                outputs = model.generate(
-                    **inputs,
-                    max_new_tokens=100,  # Reduced for CPU
-                    do_sample=True,
-                    top_k=50,
-                    top_p=0.9,
-                    temperature=0.7,
-                    pad_token_id=tokenizer.pad_token_id,
-                    eos_token_id=tokenizer.eos_token_id,
-                    use_cache=True,
-                    repetition_penalty=1.2,
-                    num_beams=1  # Disable beam search for speed
-                )
-            
-            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            
-            # Manual garbage collection
-            gc.collect()
-            
-            # Update history
-            history = history + f"\nUser: {input_text}\nAthena: {response}\n"
-            return response, history
+        
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Manual garbage collection
+        gc.collect()
+        
+        # Update history
+        history = history + f"\nUser: {input_text}\nAthena: {response}\n"
+        return response, history
             
     except Exception as e:
         return f"Error: {str(e)}", history
 
-# Load the model and tokenizer globally
-print("Loading model...")
-model, tokenizer = load_model()
-print("Model loaded successfully!")
+print("Starting app...")
 
-# Define the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("""
         # Athena-1.1.5B Chat
-        Small, Smart, Super.
+        Smaller and Smarter.
     """)
     
     # Initialize state and components
@@ -128,12 +120,5 @@ with gr.Blocks() as demo:
         outputs=[output_text, state]
     )
 
-# Launch with CPU-optimized settings
-demo.launch(
-    server_name="0.0.0.0",
-    share=False,
-    show_error=True,
-    enable_queue=True,
-    max_threads=2,  # Match your vCPU count
-    server_port=7860
-)
+if __name__ == "__main__":
+    demo.launch()