app.py

import gradio as gr
import torch
from transformers import pipeline

# --- Configuration ---
# MODEL: Using TinyLlama, a model small enough to run on a free CPU Space.
# This approach is self-contained and does not use any external API.
MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

# The "personality" of your bot.
SYSTEM_PROMPT = """You are an expert viral video scriptwriter. Your sole function is to generate compelling video hooks. When a user gives you a topic, generate a list of 10 unique and powerful video hooks. Format the output as a numbered list and do not add any extra commentary."""

# --- Model Loading ---
# We use the 'pipeline' from transformers for a simple and robust way to run the model on CPU.
# This will take a few minutes to load the first time the Space starts.
try:
    pipe = pipeline("text-generation",
                    model=MODEL_ID,
                    torch_dtype=torch.bfloat16,
                    device_map="auto")
except Exception as e:
    raise gr.Error(f"Failed to load the model. Error: {e}")

# --- The Core Chat Logic ---
def predict(message, history):
    # PROMPT FORMAT: TinyLlama uses a specific chat template.
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": message}
    ]
    
    # Use the pipeline's built-in chat template feature to format the prompt correctly.
    prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    
    # Generate the response.
    # We add a thinking status for the user because this will be slow.
    gr.Info("Generating response... this may take up to 60 seconds.")
    
    outputs = pipe(prompt,
                   max_new_tokens=1024,
                   do_sample=True,
                   temperature=0.7,
                   top_k=50,
                   top_p=0.95)
    
    # Extract the generated text from the output.
    assistant_response = outputs[0]["generated_text"].split("<|assistant|>")[-1].strip()
    return assistant_response

# --- Gradio User Interface ---
chatbot = gr.ChatInterface(
    predict,
    title="Viral Video Hook Generator",
    description="Give me a topic, and I'll generate 10 compelling video hooks. This app runs on free hardware, so please be patient with response times.",
    theme="soft",
    examples=["Productivity hacks", "The history of coffee", "How to learn a new skill"],
)

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