app.py

from unsloth import FastModel
import torch
import gc

# Set torch parameter to avoid error message, "FailOnRecompileLimitHit: recompile_limit reached with one_graph=True." when doing inference on images
torch._dynamo.config.cache_size_limit = 32   

# Initialize model
model, tokenizer = FastModel.from_pretrained(
    model_name = "unsloth/gemma-3n-E4B-it-unsloth-bnb-4bit",
#    model_name = "unsloth/gemma-3n-E2B-it", # This runs out of memory for the recommend/analyze chats
    dtype = None, # None for auto detection
    max_seq_length = 1024, # Choose any for long context!
    load_in_4bit = True,  # 4 bit quantization to reduce memory
    full_finetuning = False, # [NEW!] We have full finetuning now!
    # token = "hf_...", # use one if using gated models
)

# Helper function for inference
def do_gemma_3n_inference(model, messages, max_new_tokens = 128):
    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt = True, # Must add for generation
        tokenize = True,
        return_dict = True,
        return_tensors = "pt",
    ).to("cuda")

    with torch.no_grad():  # Disable gradient calculation during inference
        outputs = model.generate(
            **inputs,
            max_new_tokens = max_new_tokens,
            temperature = 1.0, top_p = 0.95, top_k = 64,
            return_dict_in_generate=True,  # Crucial: Get the full output
        )

    # Decode generated tokens
    outputs_excluding_inputs = outputs.sequences[:, inputs.input_ids.shape[1]:]    # exclude input tokens
    generated_text = tokenizer.batch_decode(outputs_excluding_inputs, skip_special_tokens=True)[0]

    # Cleanup to reduce VRAM usage
    del inputs
    torch.cuda.empty_cache()
    gc.collect()

    return generated_text
    
import ast

def query_ai_text_image(text, image_path=None): 
    ''' Query AI with a prompt that includes text and an image. '''
    if image_path is None:
        return "No image uploaded."
    messages = [{
        "role" : "user",
        "content": [
            { "type": "image", "image" : image_path },
            { "type": "text",  "text" : text }
        ]
    }]
    text = do_gemma_3n_inference(model, messages, max_new_tokens = 256)
    return ast.literal_eval(text)


def query_ai_text(text): 
    ''' Query AI with a text prompt. '''
    messages = [{
        "role" : "user",
        "content": [
            { "type": "text",  "text" : text }
        ]
    }]
    text = do_gemma_3n_inference(model, messages, max_new_tokens = 256)
    return text

import pandas as pd

class Inventory:
    column_names = ['title', 'author', 'year_published', 'isbn', 'description', 'copies_on_shelf', 'total_copies']

    def __init__(self, input_file_path, output_file_path): 
        ''' Initialize library inventory with data from an input csv file. Specify the file path for storing updated inventory. '''

        # Load input file, keeping only the relevant columns
        data = pd.read_csv(input_file_path)
        data = data[ [col for col in data.columns if col in self.column_names] ]

        # Check if input contains the required fields of "title" and "description"
        for col in ['title', 'description']: 
            if col not in data.columns: 
                raise Exception(f"Input book info must contain '{col}'.")
        
        # If the number of copies is not available in the input data, set it to the default value of 1
        for col in ['copies_on_shelf', 'total_copies']: 
            if col not in data.columns: 
                print(f"Input {col} not found. Setting to default value 1.")
                data[col] = 1

        # self.data = data
        # NOTE: Due to runtime memory limitations, we only demonstrate the application on the subset of books that have short descriptions.
        self.data = data[data.description.str.count(' ') < 50]
        self.file_path = output_file_path
        self.save()

    
    def save(self): 
        ''' Save inventory data to file. '''
        self.data.to_csv(self.file_path, index=False)
    
    def get_index(self, title): 
        ''' Return a pandas Index list of book(s) that match a given title. '''
        idx = self.data[self.data.title.str.lower() == title.lower()].index
        if idx.size == 0: 
            return None
        if idx.size > 1: 
            raise Exception(f"Found {idx.size} books with the title '{title}'.") #TODO: Match on author as well.
        return idx[0]

    def check_out(self, title): 
        i = self.get_index(title)
        if i is None: 
            return "ERROR: Title not found in library collection." # TODO: Add book to collection
        if self.data.loc[i, 'copies_on_shelf'] == 0: 
            return "ERROR: Check out unsuccessful. There are 0 copies on shelf."
        self.data.loc[i, 'copies_on_shelf'] -= 1
        self.save()
        return f"Check out successful. {self.data.loc[i, 'copies_on_shelf']} of {self.data.loc[i, 'total_copies']} copies remaining."

    def check_in(self, title): 
        i = self.get_index(title)
        if i is None: 
            return "ERROR: Title not found in library collection."
        row = self.data.loc[i]
        if row.copies_on_shelf == row.total_copies: 
            return f"ERROR: Check in unsuccessful. {row.copies_on_shelf} of {row.total_copies} copies already on shelf."
        self.data.loc[i, 'copies_on_shelf'] += 1
        self.save()
        return f"Check in successful. {self.data.loc[i, 'copies_on_shelf']} of {self.data.loc[i, 'total_copies']} copies on shelf."

    def get_on_shelf_book_info(self): 
        ''' Return the title/author/description info of all books with available copies on shelf, in csv format. '''
        columns = ['title', 'author', 'description']
        return self.data[self.data.copies_on_shelf > 0][columns].to_csv()

    def get_df(self): 
        ''' Return inventory data. '''
        return self.data

    def get_dtypes(self): 
        ''' Get data types for each column. '''
        return self.data.dtypes

    def set_df(self, data): 
        ''' Set inventory as the input DataFrame. '''
        self.data = data

# Initialize mobile library Inventory object
initial_book_list = '/kaggle/input/caldecott-medal-winners-1938-2019/caldecott_winners.csv'
inventory_file_path = '/kaggle/working/inventory.csv'
inventory = Inventory(initial_book_list, inventory_file_path)

import gradio as gr
from datetime import datetime


# --- "Scan" tab ---
def scan_book(image, action):

    # Query AI to extract the title and author
    prompt = "Extract the title and author from this book cover image. Format the output as ('[title]', '[author]'). If unsuccessful, output ('Unknown Title', 'Unknown Author')."
    title, author = query_ai_text_image(prompt, image)

    # AI query success check
    if title == "Unknown Title" or author == "Unknown Author":
        return "Could not reliably extract book information from the image. Please try again with a clearer cover."

    # Get the right function (check out or check in)
    if action == 'out': 
        fn = inventory.check_out
    elif action == 'in': 
        fn = inventory.check_in
    else: 
        raise Exception(f'Unknown action {action}. Valid options are "out" or "in".')

    # Perform action and return results
    return f"Title: {title}\nAuthor: {author}\n" + fn(title)


# --- "Recommend" tab ---
recommend_examples = [
    ["Suggest five books for a toddler who loves animals."],
    ["Find 3 books for a preschooler interested in space."],
    ["What are some books about adventures?"]
]

def recommend_chat_response(message, history): 
    prompt = "You are a helpful librarian making book recommendations based on the user's description of the reader's background and interests. Respond with 3-5 books, unless otherwise specified by the user. Respond with a bullet point list formatted '[title] by [author]', followed by a short sentence of less than 20 words about why this book was chosen. You must only choose books from the following csv file: " + inventory.get_on_shelf_book_info()
    return query_ai_text(f"{prompt} \n User question: {message}")


# --- "Analyze" tab ---
analyze_examples = [
    ["What is the newest book we have?"], 
    ["Summarize the common themes in our collection."]
]

def analyze_chat_response(message, history): 
    prompt = "You are a helpful librarian answering questions about the library's collection of books, based only on this inventory data: " + inventory.get_df().to_csv(index=False)
    return query_ai_text(f"{prompt} \n User question: {message}")


# --- "Manage" tab ---
def save_inventory(df_input): 
    ''' Save the user-edited DataFrame as the inventory DataFrame. ''' # TODO: More robust error checks
    df = pd.DataFrame(df_input)

    # Explicitly convert columns to desired data types
    col_type = inventory.get_dtypes().to_list()
    for i,col in enumerate(df.columns): 
        df[col] = df[col].astype(col_type[i])

    # Save DataFrame
    inventory.set_df(df)
    inventory.save()


# --- Main gradio app ---
with gr.Blocks() as demo:
    gr.Markdown("# 🚐 MoLi: Mobile Librarian 📚")
    gr.Markdown("Scan to check out/in, get book recommendations, and analyze your collection, powered by Google's Gemma 3n AI!")

    with gr.Tabs() as tabs: 

        # Scan book to check out or check in
        actions = ['out', 'in']
        with gr.Tab(label='Scan'):
            image_input = gr.Image(type='filepath', label="Upload book cover or take a photo", sources=['upload', 'webcam'], width=300)
            with gr.Row(): 
                button = {a: gr.Button(f'Check {a}') for a in actions}
            status_text = gr.Textbox(show_label=False)
            button['out'].click(fn=lambda x: scan_book(x, 'out'), inputs=image_input, outputs=status_text)
            button['in'].click(fn=lambda x: scan_book(x, 'in'), inputs=image_input, outputs=status_text)
            # # Somehow the following does not work:
            # for a, b in button.items():
            #     b.click(fn=lambda x: scan_book(x, a), inputs=image_input, outputs=status_text)

        with gr.Tab(label='Recommend'): 
            recommend_greeting = "Tell me the reader's background and interests, and I'll recommend some books available for check out!"
            gr.ChatInterface(
                fn=recommend_chat_response,
                type='messages',
                examples=recommend_examples,
                chatbot=gr.Chatbot(type='messages', placeholder=recommend_greeting),
            )
        
        with gr.Tab(label='Analyze'): 
            analyze_greeting = "Ask me anything about the library collection!"
            gr.ChatInterface(
                fn=analyze_chat_response,
                type='messages',
                examples=analyze_examples,
                chatbot=gr.Chatbot(type='messages', placeholder=analyze_greeting),
            )

        with gr.Tab(label='Manage'): 

            # Buttons
            with gr.Row(): 
                reload_button = gr.Button('Reload')
                save_button = gr.Button('Save changes')
            
            # Textbox to display status messages
            status_message = gr.Textbox(show_label=False, value='Please reload after check out or check in.')

            # Inventory table
            inventory_table = gr.DataFrame(
                value=inventory.get_df(),
                interactive=True, # Allow editing
                label="Current Library Inventory", 
                wrap=True
#                column_widths=["1fr"]*len(inventory.get_dtypes())
            )

            # Attach functions to buttons
            reload_button.click(fn=inventory.get_df, outputs=inventory_table).then(fn=lambda:f"Reloaded on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", outputs=[status_message])
            save_button.click(fn=save_inventory, inputs=inventory_table).then(fn=lambda:f"Saved on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", outputs=[status_message])


if __name__ == '__main__':
    demo.launch()