import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
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"

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, 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(
    fn=generate_response,
    inputs=gr.Textbox(label="Enter your prompt", lines=4, placeholder="Ask about Plutus..."),
    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=10).launch(enable_queue=True, max_threads=1)