app.py

# -*- coding: utf-8 -*-
"""custom-short-story-generator.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1xJz__90fcjrsG3BOWlEbHXlHhHGc8Y9I
"""

!pip install langchain>=0.1.17  transformers>=4.40.1    langchain_community gradio

!CMAKE_ARGS="-DLLAMA_CUDA=on" pip install llama-cpp-python==0.2.69

# Import necessary libraries
import gradio as gr
from llama_cpp import Llama
from langchain.llms import LlamaCpp
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

#  Initialize global variable for the model
llm = None

#  Function to load the model once
def load_model():
    global llm
    if llm is None:  # Check if the model is already loaded
        llm = LlamaCpp(
            model_path="Phi-3-mini-4k-instruct-fp16.gguf",
            n_gpu_layers=-1,  # Adjusted for memory efficiency
            max_tokens=500,
            n_ctx=1024,  # Lower memory usage
            seed=42,
            verbose=False
        )

#  Define the story prompts for each type
story_prompts = {
    "Classic Short Story": """
    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}.
    """,
    "Suspense and Conflict Focus": """
    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}.
    """,
    "Character-Driven Story": """
    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}.
    """,
    "World-Building and Immersion": """
    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}.
    """,
    "Emotional and Impactful Narrative": """
    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}.
    """
}

#  Function to generate the story based on user input and selected story type
def generate_story(paragraph, story_type):
    load_model()  # Load the model if it's not already loaded
    prompt_template = story_prompts[story_type]
    prompt = prompt_template.format(paragraph=paragraph)

    # Generate the story using the loaded model
    output = llm.invoke(
        input=prompt  # Use the formatted prompt
    )

    # Clean the output (remove <|assistant|> and ensure the story has proper ending)
    story = output.replace("<|assistant|>", "").strip()
    if not story.endswith((".", "!", "?")):
        story += "."

    return story

# Define the Gradio interface
def story_interface(paragraph, story_type):
    generated_story = generate_story(paragraph, story_type)
    return generated_story

#  Define available story types for dropdown
story_types = [
    "Classic Short Story",
    "Suspense and Conflict Focus",
    "Character-Driven Story",
    "World-Building and Immersion",
    "Emotional and Impactful Narrative"
]

#  Set up Gradio interface with enhanced design
gr_interface = gr.Interface(
    fn=story_interface,  # Function that generates the story
    inputs=[
        gr.Textbox(label="Enter your idea (1-3 sentences)", placeholder="A young detective embarks on a mission to uncover the truth..."),
        gr.Dropdown(story_types, label="Select Story Type")
    ],
    outputs=gr.Textbox(label="Generated Story"),  # Display the generated story
    title="Custom Short Story Generator",
    description="Create captivating stories based on your input! Simply enter a short paragraph and select a story type to generate a unique, complete narrative.",
    theme="default",  #  change theme (can be 'default', 'dark', etc.)
    css=".interface-title { font-size: 2em; color: #2c3e50; }"  # Add custom styling
)

# Launch the Gradio app
gr_interface.launch()