Update app.py

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import spaces  # Explicit import for GPU decorator
+import spaces
 import logging
+import json
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -12,40 +13,44 @@ logger = logging.getLogger(__name__)
 model = None
 tokenizer = None
 MODEL_NAME = "ubiodee/Test_Plutus"
-FALLBACK_TOKENIZER = "gpt2"
+FALLBACK_TOKENIZER = "NousResearch/Meta-Llama-3-8B"
 
 # Load tokenizer at startup (lightweight, no model yet)
 try:
-    logger.info("Loading tokenizer at startup with legacy versions...")
+    logger.info("Loading tokenizer at startup for %s...", MODEL_NAME)
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
-        use_fast=False,
+        use_fast=True,  # Llama-3 uses fast tokenizer
         trust_remote_code=True,
     )
     logger.info("Primary tokenizer loaded successfully.")
 except Exception as e:
-    logger.warning(f"Primary tokenizer failed: {str(e)}. Using fallback.")
-    tokenizer = AutoTokenizer.from_pretrained(
-        FALLBACK_TOKENIZER,
-        use_fast=False,
-        trust_remote_code=True,
-    )
-    logger.info("Fallback tokenizer loaded.")
+    logger.warning(f"Primary tokenizer failed: {str(e)}. Using fallback: {FALLBACK_TOKENIZER}")
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(
+            FALLBACK_TOKENIZER,
+            use_fast=True,
+            trust_remote_code=True,
+        )
+        logger.info("Fallback tokenizer loaded successfully.")
+    except Exception as fallback_e:
+        logger.error(f"Fallback tokenizer failed: {str(fallback_e)}")
+        raise
 
 # Set pad token
 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.")
+    logger.info("Set pad_token_id to eos_token_id: %s", tokenizer.eos_token_id)
 
 def load_model():
     """Load model inside GPU context."""
     global model
     if model is None:
         try:
-            logger.info("Loading model with CPU fallback (full precision)...")
+            logger.info("Loading model %s with torch.float16...", MODEL_NAME)
             model = AutoModelForCausalLM.from_pretrained(
                 MODEL_NAME,
-                torch_dtype=torch.float16,  # Use fp16 for memory efficiency
+                torch_dtype=torch.float16,  # Use fp16 for ZeroGPU
                 low_cpu_mem_usage=True,
                 trust_remote_code=True,
             )
@@ -69,7 +74,7 @@ def generate_response(prompt, progress=gr.Progress()):
     
     progress(0.3, desc="Tokenizing input...")
     try:
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
         
         progress(0.6, desc="Generating response...")
         with torch.no_grad():
@@ -97,11 +102,11 @@ def generate_response(prompt, progress=gr.Progress()):
 # Gradio UI
 demo = gr.Interface(
     fn=generate_response,
-    inputs=gr.Textbox(label="Enter your prompt", lines=4, placeholder="Ask about Plutus..."),
+    inputs=gr.Textbox(label="Enter your prompt", lines=4, placeholder="Ask about Plutus smart contracts..."),
     outputs=gr.Textbox(label="Model Response"),
     title="Cardano Plutus AI Assistant",
     description="Write Plutus smart contracts on Cardano blockchain."
 )
 
-# Launch with queueing
-demo.queue(max_size=5).launch(enable_queue=True, max_threads=1)
+# Launch with simplified queueing
+demo.launch(queue=True, max_queue_size=5)