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

# --- Configuration ---
# Your custom model from the Hugging Face Hub
MODEL_NAME = "vignesh-ramesh/myemoji-gemma-3-270m-it" 

# Load the model and tokenizer. 
# device_map="auto" handles placing the model across available devices (CPU/GPU)
# This is required for models loaded with AutoModelForCausalLM.
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Create the text-generation pipeline.
# CRITICAL FIX: The 'device' argument is intentionally omitted here 
# because 'device_map="auto"' already handles device assignment.
model_pipeline = pipeline(
    "text-generation", 
    model=model, 
    tokenizer=tokenizer,
    # REMOVED: device=device
)
# ---------------------

def generate_emojis(text: str) -> str:
    """
    Predicts the emoji response for a given text prompt using the model's chat template.
    """
    # 1. Define the conversation messages with the required system instruction
    inference_messages = [
        # System instruction is crucial for guiding the model's behavior
        {"role": "system", "content": "Translate this text to emoji: "},
        {"role": "user", "content": text}
    ]
    
    # 2. Apply the model's specific chat template (Gemma format) to the prompt
    prompt = tokenizer.apply_chat_template(
        inference_messages, 
        tokenize=False, 
        add_generation_prompt=True
    )

    # 3. Generate the output
    response = model_pipeline(
        prompt,
        max_new_tokens=15,          # Limit generation length for emojis
        do_sample=True,             # Use sampling for variety
        temperature=0.7,            # Set temperature for creativity
        return_full_text=True,      # Get full output to strip the prompt
        num_return_sequences=1
    )
    
    # 4. Extract and clean the generated text
    full_output = response[0]['generated_text']
    
    # Strip the original prompt from the full output to get only the generation
    emojis_text = full_output[len(prompt):].strip()
        
    return emojis_text

# Create the Gradio Interface
demo = gr.Interface(
    fn=generate_emojis,
    inputs=gr.Textbox(
        lines=3, 
        label="Enter Text for Emoji Generation", 
        placeholder="e.g., I just aced my final exam!"
    ),
    outputs=gr.Textbox(
        label="Equivalent Emojis", 
        type="text",
        show_copy_button=True
    ),
    title="Text to Emoji Generator (Gemma 3.1 270M)",
    description="Input text and get the equivalent emojis using a fine-tuned Gemma model.",
    allow_flagging='never',
    examples=[
        "This is the best day ever.",
        "A beautiful sunset by the beach",
        "It's time for coffee and a big slice of cake."
    ]
)

# Launch the app
if __name__ == "__main__":
    demo.launch()