custom_short_story_generator.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PpZ2FLMBU5FL"
      },
      "outputs": [],
      "source": [
        "!pip install langchain>=0.1.17  transformers>=4.40.1    langchain_community gradio"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!CMAKE_ARGS=\"-DLLAMA_CUDA=on\" pip install llama-cpp-python==0.2.69"
      ],
      "metadata": {
        "id": "Jh_DeTdj4AL4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Import necessary libraries\n",
        "import gradio as gr\n",
        "from llama_cpp import Llama\n",
        "from langchain.llms import LlamaCpp\n",
        "from langchain.prompts import PromptTemplate\n",
        "from langchain.chains import LLMChain"
      ],
      "metadata": {
        "id": "U2a_q--S4ChK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#  Initialize global variable for the model\n",
        "llm = None\n",
        "\n",
        "#  Function to load the model once\n",
        "def load_model():\n",
        "    global llm\n",
        "    if llm is None:  # Check if the model is already loaded\n",
        "        llm = LlamaCpp(\n",
        "            model_path=\"Phi-3-mini-4k-instruct-fp16.gguf\",\n",
        "            n_gpu_layers=-1,  # Adjusted for memory efficiency\n",
        "            max_tokens=500,\n",
        "            n_ctx=1024,  # Lower memory usage\n",
        "            seed=42,\n",
        "            verbose=False\n",
        "        )"
      ],
      "metadata": {
        "id": "pIiiCTzT4IIO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#  Define the story prompts for each type\n",
        "story_prompts = {\n",
        "    \"Classic Short Story\": \"\"\"\n",
        "    Using this paragraph, write a classic short story with well-rounded characters and vivid settings. Ensure the story has a beginning, middle, and end with a clear conflict driving the narrative. Focus on character depth and concise storytelling. Complete the story in full. User's paragraph: {paragraph}.\n",
        "    \"\"\",\n",
        "    \"Suspense and Conflict Focus\": \"\"\"\n",
        "    Create a suspenseful story from this paragraph. Build tension using conflict, withholding information, and stakes that grow as the story progresses. Ensure the story is fully resolved and not cut off. User's paragraph: {paragraph}.\n",
        "    \"\"\",\n",
        "    \"Character-Driven Story\": \"\"\"\n",
        "    Write a character-driven story based on this paragraph. Focus on character development, growth, and internal conflict. Ensure a clear beginning, middle, and end, completing the story fully. User's paragraph: {paragraph}.\n",
        "    \"\"\",\n",
        "    \"World-Building and Immersion\": \"\"\"\n",
        "    Based on this paragraph, craft a story rich in world-building and immersion. Focus on creating a detailed setting, ensuring a full plot with clear beginning, middle, and end. Complete the story without cutting off. User's paragraph: {paragraph}.\n",
        "    \"\"\",\n",
        "    \"Emotional and Impactful Narrative\": \"\"\"\n",
        "    Create an emotionally impactful story based on this paragraph. Focus on character emotions and stakes that resonate deeply. Ensure the story has a full arc, with a clear beginning, middle, and end. User's paragraph: {paragraph}.\n",
        "    \"\"\"\n",
        "}"
      ],
      "metadata": {
        "id": "iBPIvtmc4P9F"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#  Function to generate the story based on user input and selected story type\n",
        "def generate_story(paragraph, story_type):\n",
        "    load_model()  # Load the model if it's not already loaded\n",
        "    prompt_template = story_prompts[story_type]\n",
        "    prompt = prompt_template.format(paragraph=paragraph)\n",
        "\n",
        "    # Generate the story using the loaded model\n",
        "    output = llm.invoke(\n",
        "        input=prompt  # Use the formatted prompt\n",
        "    )\n",
        "\n",
        "    # Clean the output (remove <|assistant|> and ensure the story has proper ending)\n",
        "    story = output.replace(\"<|assistant|>\", \"\").strip()\n",
        "    if not story.endswith((\".\", \"!\", \"?\")):\n",
        "        story += \".\"\n",
        "\n",
        "    return story"
      ],
      "metadata": {
        "id": "PYV30lz54V8l"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the Gradio interface\n",
        "def story_interface(paragraph, story_type):\n",
        "    generated_story = generate_story(paragraph, story_type)\n",
        "    return generated_story"
      ],
      "metadata": {
        "id": "JLYTR7DC4cb8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#  Define available story types for dropdown\n",
        "story_types = [\n",
        "    \"Classic Short Story\",\n",
        "    \"Suspense and Conflict Focus\",\n",
        "    \"Character-Driven Story\",\n",
        "    \"World-Building and Immersion\",\n",
        "    \"Emotional and Impactful Narrative\"\n",
        "]"
      ],
      "metadata": {
        "id": "V6x9ekjX4n23"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#  Set up Gradio interface with enhanced design\n",
        "gr_interface = gr.Interface(\n",
        "    fn=story_interface,  # Function that generates the story\n",
        "    inputs=[\n",
        "        gr.Textbox(label=\"Enter your idea (1-3 sentences)\", placeholder=\"A young detective embarks on a mission to uncover the truth...\"),\n",
        "        gr.Dropdown(story_types, label=\"Select Story Type\")\n",
        "    ],\n",
        "    outputs=gr.Textbox(label=\"Generated Story\"),  # Display the generated story\n",
        "    title=\"Custom Short Story Generator\",\n",
        "    description=\"Create captivating stories based on your input! Simply enter a short paragraph and select a story type to generate a unique, complete narrative.\",\n",
        "    theme=\"default\",  #  change theme (can be 'default', 'dark', etc.)\n",
        "    css=\".interface-title { font-size: 2em; color: #2c3e50; }\"  # Add custom styling\n",
        ")"
      ],
      "metadata": {
        "id": "LDZ9hpa_4n02"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Launch the Gradio app\n",
        "gr_interface.launch()"
      ],
      "metadata": {
        "id": "babiTcI_4wfy"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}