import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Check for GPU and use CPU if not available (your zero-GPU approach)
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# Load model and tokenizer
model_id = "google/gemma-2-2b-it"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16  # Uses less memory
)

def chat_response(message, history):
    # Format the prompt according to Gemma 2's template
    prompt = f"<start_of_turn>user\n{message}<end_of_turn>\n<start_of_turn>model"
    
    # Tokenize input
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    # Generate response
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    
    # Decode the response
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Extract only the model's response
    if "<start_of_turn>model" in response:
        return response.split("<start_of_turn>model")[-1].strip()
    return response

# Create the chat interface
demo = gr.ChatInterface(
    fn=chat_response,
    title="Hausa AI Assistant",
    description="A simple AI assistant powered by Gemma 2 2B"
)

# Launch the application
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)