app.py

import gradio as gr
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("facebook/galactica-1.3b")
model = AutoModelForCausalLM.from_pretrained("facebook/galactica-1.3b")
text2text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer, num_workers=2)

def predict(text, max_length=64, penalty_alpha=0.6, top_k=4):
    text = text.strip()
    out_text = text2text_generator(text, max_length=max_length, 
                              penalty_alpha=penalty_alpha, 
                              top_k=top_k,
                              eos_token_id = tokenizer.eos_token_id,
                              bos_token_id = tokenizer.bos_token_id,
                              pad_token_id = tokenizer.pad_token_id,
                         )[0]['generated_text']
    out_text = "<p>" + out_text + "</p>"
    out_text = out_text.replace(text, text + "<b><span>")
    out_text = out_text +  "</span></b>"
    out_text = out_text.replace("\n", "<br>")
    return out_text

iface = gr.Interface(
    fn=predict, 
    inputs=[
        gr.inputs.Textbox(lines=5, label="Input Text"),
        gr.inputs.Slider(minimum=32, maximum=64, default=64, label="Max Length"),
        gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.6, step=0.1, label="Penalty Alpha"),
        # gr.inputs.Checkbox(label="Do Sample"),
        gr.inputs.Slider(minimum=0, maximum=16, default=8, step=1, label="Top K")
    ],
    outputs=gr.HTML(),
    description="Galactica Base Model",
    examples=[[
            "The attention mechanism in LLM is",
            32,
            0.6,
            4
        ], 
        [
            "Title: Attention is all you need\n\nAbstract:",
            32,
            0.6,
            4
        ]
    ]
)

iface.launch()