app.py

import gradio as gr
import torch
import torch.multiprocessing as mp
from transformers import AutoModelForCausalLM, AutoTokenizer
import spaces
import logging

# 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
MODEL_NAME = "ubiodee/plutus_llm"

# 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.")

# Gradio UI
demo = gr.Interface(
    fn=generate_response,
    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
demo.launch()