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

# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-3B")

def predict(input_text):
    """Generate a response using the Llama model."""
    inputs = tokenizer.encode(input_text, return_tensors="pt")
    outputs = model.generate(inputs, max_length=150, num_return_sequences=1)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Metadata
title = "Llama 3.2-3B Model Demonstration"
description = """
## How to Use
1. Enter a prompt in the input box.
2. Click 'Submit' to get the model's response.
3. Explore the examples provided for inspiration.

## Model Details
- **Name:** Llama 3.2-3B
- **Capabilities:** Text generation, summarization, translation, etc.

This Space demonstrates the capabilities of Meta's Llama 3.2-3B model.
"""
examples = [
    ["Generate a summary for: Artificial intelligence is the simulation of human intelligence..."],
    ["Translate the text: Hello, how are you? into French."],
    ["What are the benefits of renewable energy sources?"],
    ["Write a poem about the ocean."],
]

# Interface
def create_interface():
    """Create a Gradio interface for the Llama model."""
    return gr.Interface(
        fn=predict,
        inputs=gr.Textbox(lines=5, placeholder="Type your input here...", label="Enter your prompt"),
        outputs=gr.Textbox(label="Model Output"),
        title=title,
        description=description,
        examples=examples,
        theme="compact"
    )

# Launch the interface
interface = create_interface()
interface.launch(share=True, server_name="0.0.0.0")