app.py

import os
import torch
import gradio as gr
import spaces
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer


# Use a global variable to hold the current model and tokenizer
current_model = None
current_tokenizer = None

def load_model_on_selection(model_name, progress=gr.Progress(track_tqdm=False)):
    global current_model, current_tokenizer
    token = os.getenv("HF_TOKEN")

    progress(0, desc="Loading tokenizer...")
    current_tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)

    progress(0.5, desc="Loading model...")
    current_model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map="cpu",
        use_auth_token=token
    )

    progress(1, desc="Model ready.")
    return f"{model_name} loaded and ready!"

@spaces.GPU
def generate_text(prompt):
    global current_model, current_tokenizer
    if current_model is None or current_tokenizer is None:
        yield "⚠️ No model loaded yet. Please select a model first."

    current_model.to("cuda")
    inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)

    output_ids = []
    streamer_output = ""

    for token_id in current_model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=False,
        return_dict_in_generate=True,
        output_scores=False
        ).sequences[0]:
        
        output_ids.append(token_id.item())

        yield current_tokenizer.decode(output_ids, skip_special_tokens=True)


# Model options
model_choices = [
    "meta-llama/Llama-3.2-3B-Instruct",
    "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
    "google/gemma-7b"
]

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## Clinical Text Testing with LLaMA, DeepSeek, and Gemma")

    model_selector = gr.Dropdown(choices=model_choices, label="Select Model")
    model_status = gr.Textbox(label="Model Status", interactive=False)

    input_text = gr.Textbox(label="Input Clinical Text")
    generate_btn = gr.Button("Generate")

    output_text = gr.Textbox(label="Generated Output")

    # Load model on dropdown change
    model_selector.change(fn=load_model_on_selection, inputs=model_selector, outputs=model_status)

    # Generate with current model
    generate_btn.click(fn=generate_text, inputs=input_text, outputs=output_text)
    input_text.submit(fn=generate_text, inputs=input_text, outputs=output_text)


load_model_on_selection("meta-llama/Llama-3.2-3B-Instruct")
demo.launch()