import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel, PeftConfig
import torch
import logging
import gc

def clear_memory():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

def load_model():
    # Force CPU for stability
    device = torch.device("cpu")
    
    # Load base model
    base_model = AutoModelForCausalLM.from_pretrained(
        "microsoft/phi-2",
        torch_dtype=torch.float32,
        trust_remote_code=True,
        device_map=None,
        low_cpu_mem_usage=True
    )
    
    # Load LoRA configuration
    peft_config = PeftConfig.from_pretrained("phi2-oasst-qlora-final")
    
    # Load and merge LoRA weights
    model = PeftModel.from_pretrained(
        base_model,
        "phi2-oasst-qlora-final",
        device_map=None,
        torch_dtype=torch.float32
    )
    
    # Merge LoRA weights with base model
    model = model.merge_and_unload()
    
    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        "microsoft/phi-2",
        trust_remote_code=True
    )
    tokenizer.pad_token = tokenizer.eos_token
    
    return model, tokenizer

def generate_response(prompt):
    # Format prompt
    full_prompt = f"Instruct: {prompt}\nOutput:"
    
    # Tokenize
    inputs = tokenizer(
        full_prompt, 
        return_tensors="pt",
        padding=False,
        truncation=True,
        max_length=256
    )
    
    # Generate
    with torch.inference_mode():
        outputs = model.generate(
            **inputs,
            max_length=128,
            min_length=10,
            temperature=0.7,
            top_p=0.9,
            num_return_sequences=1,
            pad_token_id=tokenizer.eos_token_id,
            do_sample=True,
            no_repeat_ngram_size=2,
            use_cache=True
        )
    
    # Decode response
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Remove the prompt from response
    response = response.replace(full_prompt, "").strip()
    
    return response

# Load model and tokenizer globally
print("Loading model...")
model, tokenizer = load_model()
print("Model loaded!")

# Create Gradio interface
iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(
        lines=3, 
        placeholder="Enter your instruction here...",
        label="Input"
    ),
    outputs=gr.Textbox(
        lines=5,
        label="Generated Response"
    ),
    title="Phi-2 Fine-tuned Assistant",
    description="A fine-tuned version of Phi-2 on the OpenAssistant dataset",
    examples=[
        ["Explain what is Python in one sentence."],
        ["How do neural networks learn through backpropagation?"],
        ["Write a short poem about coding."]
    ]
)

if __name__ == "__main__":
    iface.launch()