import spaces
import gradio as gr
import torch

BASE_MODEL = "deepseek-ai/deepseek-coder-1.3b-instruct"
ADAPTER = "praveends/migration-copilot-deepseek-coder-1-3b-instruct"

# Global variables — loaded lazily inside GPU function
tokenizer = None
model = None

def load_model():
    global tokenizer, model
    if model is not None:
        return
    from transformers import AutoTokenizer, AutoModelForCausalLM
    from peft import PeftModel
    print("Loading model...")
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.float16,
        trust_remote_code=True,
        device_map="auto",
    )
    model = PeftModel.from_pretrained(model, ADAPTER)
    model = model.merge_and_unload()
    model.eval()
    print("Model loaded!")

@spaces.GPU(duration=60)
def generate(prompt: str) -> str:
    load_model()
    inputs = tokenizer(
        prompt, return_tensors="pt",
        truncation=True, max_length=512
    ).to("cuda")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=300,
            temperature=0.1,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response[len(prompt):].strip()

demo = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(label="Prompt", lines=10),
    outputs=gr.Textbox(label="Generated PySpark", lines=10),
    title="Migration Copilot Inference",
    api_name="generate",
)

demo.launch()