Update app.py

app.py

@@ -1,44 +1,87 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from spaces import GPU  # Import for GPU decorator
+from spaces import GPU
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Load model & tokenizer
 MODEL_NAME = "ubiodee/Test_Plutus"
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
-model.eval()
+try:
+    logger.info("Loading tokenizer with use_fast=False...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_NAME,
+        use_fast=False,  # Use slow tokenizer to avoid fast tokenizer errors
+        use_safetensors=True,
+        trust_remote_code=True,  # Allow custom tokenizer code
+    )
+    logger.info("Tokenizer loaded successfully.")
+except Exception as e:
+    logger.error(f"Tokenizer loading failed: {str(e)}")
+    raise
+
+try:
+    logger.info("Loading model with 8-bit quantization...")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        device_map="auto",  # Automatically map to GPU/CPU
+        load_in_8bit=True,  # Use 8-bit quantization to match model
+        torch_dtype=torch.bfloat16,  # Use bfloat16 for efficiency
+        use_safetensors=True,
+        low_cpu_mem_usage=True,  # Reduce CPU memory during loading
+        trust_remote_code=True,  # Allow custom model code
+    )
+    model.eval()
+    logger.info("Model loaded successfully.")
+except Exception as e:
+    logger.error(f"Model loading failed: {str(e)}")
+    raise
 
 # Set pad token if not defined
 if tokenizer.pad_token_id is None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
+    logger.info("Set pad_token_id to eos_token_id.")
 
+# Move model to GPU if available
 if torch.cuda.is_available():
     model.to("cuda")
+    logger.info("Model moved to GPU.")
+else:
+    logger.warning("No GPU available, using CPU.")
 
 # Response function with GPU decorator
 @spaces.GPU
-def generate_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=200,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=tokenizer.pad_token_id,
-        )
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Remove the prompt from the output to return only the answer
-    if response.startswith(prompt):
-        response = response[len(prompt):].strip()
-
-    return response
+def generate_response(prompt, progress=gr.Progress()):
+    progress(0.1, desc="Tokenizing input...")
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        
+        progress(0.5, desc="Generating response...")
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=200,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                eos_token_id=tokenizer.eos_token_id,
+                pad_token_id=tokenizer.pad_token_id,
+            )
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Remove the prompt from the output
+        if response.startswith(prompt):
+            response = response[len(prompt):].strip()
+        
+        progress(1.0, desc="Done!")
+        return response
+    except Exception as e:
+        logger.error(f"Inference failed: {str(e)}")
+        return f"Error during generation: {str(e)}"
 
 # Gradio UI
 demo = gr.Interface(
@@ -49,5 +92,5 @@ demo = gr.Interface(
     description="Write Plutus smart contracts on Cardano blockchain."
 )
 
-# Launch with queueing for Spaces
-demo.queue().launch(enable_queue=True, max_threads=1)
+# Launch with queueing
+demo.queue(max_size=10).launch(enable_queue=True, max_threads=1)