import gradio as gr
from ctransformers import AutoModelForCausalLM

# Load the TinyLlama model with ctransformers
llm = AutoModelForCausalLM.from_pretrained(
    "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf",
    model_file="tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf",
    model_type="tinyllama",
    max_new_tokens=512
)

# Define a function to generate text based on user input
def generate_text(prompt):
    # Generate response from the model
    return llm(prompt)

# Set up Gradio interface
interface = gr.Interface(
    fn=generate_text,  # Function to call
    inputs="text",     # Text input for prompt
    outputs="text",    # Text output for response
    title="TinyLlama GGUF Text Generator",
    description="Enter a prompt and see how TinyLlama responds."
)

# Launch the app
interface.launch()