Update app.py

app.py

@@ -1,106 +1,56 @@
 import gradio as gr
 import torch
-import torch.multiprocessing as mp
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import spaces
-import logging
+from spaces import GPU  # Import ZeroGPU decorator
 
-# Set multiprocessing to 'spawn' for ZeroGPU compatibility
-try:
-    mp.set_start_method('spawn', force=True)
-except RuntimeError:
-    pass
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Global variables
-model = None
-tokenizer = None
+# Load model & tokenizer (runs on CPU at startup)
 MODEL_NAME = "ubiodee/plutus_llm"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    load_in_8bit=True
+)
 
-# Load tokenizer at startup
-try:
-    logger.info("Loading tokenizer at startup for %s...", MODEL_NAME)
-    tokenizer = AutoTokenizer.from_pretrained(
-        MODEL_NAME,
-        use_fast=True,
-        trust_remote_code=True,
-    )
-    logger.info("Primary tokenizer loaded successfully.")
-except Exception as e:
-    logger.error(f"Tokenizer loading failed: {str(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: %s", tokenizer.eos_token_id)
-
-# Load model at startup
-try:
-    logger.info("Loading model %s with torch.float16...", MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        torch_dtype=torch.float16,
-        trust_remote_code=True,
-    )
-    model.eval()
-    logger.info("Model loaded successfully.")
-except Exception as e:
-    logger.error(f"Model loading failed: {str(e)}")
-    raise
-
-# Response function
-@spaces.GPU(duration=120)
-def generate_response(prompt, progress=gr.Progress()):
-    global model
-    progress(0.1, desc="Moving model to GPU...")
-    try:
-        if torch.cuda.is_available():
-            model = model.to("cuda")
-            logger.info("Model moved to GPU.")
-        else:
-            logger.warning("GPU not available; using CPU.")
-        
-        progress(0.3, desc="Tokenizing input...")
-        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():
-            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)
-        
-        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)}"
-    finally:
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            logger.info("GPU memory cleared.")
+# Set padding token
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+model.eval()
+
+# Response function with ZeroGPU decorator
+@GPU
+def generate_response(prompt, max_new_tokens=200, temperature=0.7, top_p=0.9):
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to("cuda")
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            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)
+    if response.startswith(prompt):
+        response = response[len(prompt):].strip()
+    return response
 
 # Gradio UI
 demo = gr.Interface(
     fn=generate_response,
-    inputs=gr.Textbox(label="Enter your prompt", lines=4, placeholder="Ask about Plutus smart contracts..."),
+    inputs=[
+        gr.Textbox(label="Enter your prompt", lines=4, placeholder="Ask about Plutus..."),
+        gr.Slider(label="Max New Tokens", minimum=50, maximum=500, value=200, step=10),
+        gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, value=0.7, step=0.1),
+        gr.Slider(label="Top P", minimum=0.1, maximum=1.0, value=0.9, step=0.05)
+    ],
     outputs=gr.Textbox(label="Model Response"),
     title="Cardano Plutus AI Assistant",
-    description="Write Plutus smart contracts on Cardano blockchain."
+    description="Ask questions about Plutus smart contracts or Cardano blockchain using ubiodee/plutus_llm."
 )
 
-# Launch
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()