# app.py

import os

import gradio as gr

import keras

import keras_hub

from huggingface_hub import from_pretrained_keras



# Set Keras backend

os.environ["KERAS_BACKEND"] = "jax"

os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "1.00"



# --- 1. LOAD THE MERGED MODEL FROM THE HUB ---

# !!! IMPORTANT: Change this to your username and repo name !!!

repo_id = "Tarive/lora_research_abstracts" 

print(f"Loading merged model from Hub: {repo_id}")

gemma_lm = from_pretrained_keras(repo_id)



# Compile the model with a sampler for generation

gemma_lm.compile(sampler=keras_hub.samplers.TopKSampler(k=5))

print("Model loaded and compiled successfully.")





# --- 2. DEFINE THE INFERENCE FUNCTION ---

def revise_abstract(draft_abstract, grant_type):

    if not draft_abstract or not grant_type:

        return "Error: Please provide both a draft abstract and a grant type."



    template = (

        "Instruction:\n"

        "You are an expert grant writer. Rewrite the following draft abstract to be more impactful and clear, "

        "following the specific conventions of a {activity_code} grant. Ensure the most compelling claims are front-loaded.\n\n"

        "Input Draft:\n"

        "{unoptimized_abstract}\n\n"

        "Revised Abstract:"

    )

    prompt = template.format(unoptimized_abstract=draft_abstract, activity_code=grant_type)

    output = gemma_lm.generate(prompt, max_length=1024)

    

    parts = output.split("Revised Abstract:")

    return parts[1].strip() if len(parts) > 1 else output.strip()



# --- 3. CREATE THE GRADIO INTERFACE ---

with gr.Blocks(theme=gr.themes.Soft()) as demo:

    gr.Markdown("# Grant Abstract Revision Tool (Fine-Tuned on Gemma)")

    gr.Markdown("Enter a draft abstract and select its grant type. The model will rewrite it to be more impactful, based on patterns from successfully funded NIH grants.")

    

    with gr.Row():

        draft_input = gr.Textbox(lines=15, label="Input Draft Abstract", placeholder="Paste your draft abstract here...")

        grant_type_input = gr.Dropdown(

            ["R01", "R21", "F32", "T32", "P30", "R41", "R43", "R44", "K99"], 

            label="Grant Type (Activity Code)",

            info="Select the grant mechanism you are targeting."

        )

    

    submit_button = gr.Button("Revise Abstract", variant="primary")

    revised_output = gr.Textbox(lines=15, label="Model's Revised Abstract", interactive=False)

    

    submit_button.click(fn=revise_abstract, inputs=[draft_input, grant_type_input], outputs=revised_output)

    

    gr.Examples(

        examples=[

            ["SUMMARY \nA pressing concern exists regarding lead poisoning...This study aimed to optimize and validate a dried blood spot collection device...", "R41"],

            ["This project is about figuring out how macrophages and S. flexneri interact...Our study will look at this in vitro and in vivo...", "R21"]

        ],

        inputs=[draft_input, grant_type_input]

    )



print("Launching Gradio app...")

demo.launch()