Spaces:
Sleeping
Sleeping
File size: 2,840 Bytes
cdb4f63 97ff2e1 cdb4f63 97ff2e1 cdb4f63 97ff2e1 cdb4f63 97ff2e1 cdb4f63 97ff2e1 cdb4f63 97ff2e1 cdb4f63 97ff2e1 5e4de6b cdb4f63 97ff2e1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """
Module for loading a LoRA fine-tuned BART model and serving
an interactive Gradio interface for text generation.
"""
import torch
import gradio as gr
from transformers import AutoTokenizer
from transformers import BartForConditionalGeneration
from peft import PeftModel
def load_model() -> tuple[AutoTokenizer, PeftModel, torch.device]:
"""
Load tokenizer and LoRA-enhanced model onto available device.
Returns:
tokenizer (AutoTokenizer): Tokenizer for text processing.
model (PeftModel): Fine-tuned LoRA BART model in eval mode.
device (torch.device): Computation device (GPU if available, else CPU).
"""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load tokenizer and base model
tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
base_model = BartForConditionalGeneration.from_pretrained("facebook/bart-base")
# Use efficient attention if desired
base_model.config.attn_implementation = "sdpa"
# Load PEFT (LoRA) model for inference
model = PeftModel.from_pretrained(
base_model, "outputs/bart-base-reddit-lora"
).eval()
model.to(device)
model.eval()
return tokenizer, model, device
# Load once at startup
tokenizer, model, device = load_model()
def predict(text: str) -> str:
"""
Generate a text response given an input prompt.
Args:
text (str): The input prompt string.
Returns:
str: The decoded model output.
"""
# Tokenize and move inputs to device
inputs = tokenizer(
text,
return_tensors="pt",
padding=True,
truncation=True,
).to(device)
# Generate with both beam search and sampling for diversity
outputs = model.generate(
**inputs,
max_length=128,
num_beams=10,
do_sample=True,
length_penalty=1.2,
repetition_penalty=1.3,
no_repeat_ngram_size=3,
top_p=0.9,
temperature=0.8,
early_stopping=True,
eos_token_id=tokenizer.eos_token_id,
)
# Decode the first generated sequence
return tokenizer.decode(outputs[0], skip_special_tokens=True)
def main() -> None:
"""
Launch Gradio web interface for interactive model inference.
"""
interface = gr.Interface(
fn=predict,
inputs=gr.Textbox(lines=5, placeholder="Broad questions often have better results (e.g. What do you think about politics right now?).", label="Your Question"),
outputs=gr.Textbox(label="Mimic Bot's Comment"),
title="Reddit-User-Mimic-Bot Inference (Bart-LoRA)",
description="Enter a question you would ask on reddit, and our Mimic Bot would comment back! Have fun.",
allow_flagging="never",
)
interface.launch()
if __name__ == "__main__":
main()
|