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

# Check if CUDA is available and set the device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Load model and tokenizer
MODEL_PATH = "sagar007/phi3.5_finetune"
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

base_model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Phi-3.5-mini-instruct",
    torch_dtype=torch.float16 if device.type == "cuda" else torch.float32,
    device_map="auto",
    trust_remote_code=True
)

peft_config = PeftConfig.from_pretrained(MODEL_PATH)
model = PeftModel.from_pretrained(base_model, MODEL_PATH)
model.to(device)
model.eval()

@spaces.GPU(duration=60)
def generate_response(instruction, max_length=512):
    prompt = f"Instruction: {instruction}\nResponse:"
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=max_length,
            num_return_sequences=1,
            temperature=0.7,
            top_p=0.9,
            do_sample=True
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response.split("Response:")[1].strip()

def chatbot(message, history):
    response = generate_response(message)
    return response

demo = gr.ChatInterface(
    chatbot,
    title="Fine-tuned Phi-3.5 Chatbot",
    description="This is a chatbot using a fine-tuned version of the Phi-3.5 mini model.",
    theme="default",
    examples=[
        "Explain the concept of machine learning.",
        "Write a short story about a robot learning to paint.",
        "What are some effective ways to reduce stress?",
        "Summarize the key points of climate change in simple terms.",
        "Create a step-by-step guide for making a perfect omelette.",
        "Describe the differences between classical and quantum computing.",
        "Write a motivational speech for a team starting a new project.",
        "Explain the importance of biodiversity in ecosystems.",
        "Compose a haiku about artificial intelligence.",
        "List five tips for effective time management.",
        "Describe the process of photosynthesis in layman's terms.",
        "Write a dialogue between two characters discussing the future of space exploration.",
        "Explain the concept of blockchain technology and its potential applications."
        ],
    cache_examples=False,
)

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