app.py

import gradio as gr
from transformers import pipeline
from transformers import BloomTokenizerFast, BloomForCausalLM
import re

description = """
<img src="https://huggingface.co/spaces/tomrb/bettercallbloom/resolve/main/img.jpeg" width=300px style="margin:auto;">
When in legal doubt, you better call BLOOM! Ask BLOOM any legal question. \n
***Advice here is for informational purposes only and should not be considered final or official advice. See a local attorney for the best answer to your questions.***
"""
title = "Better Call Bloom!"



tokenizer = BloomTokenizerFast.from_pretrained("tomrb/bettercallbloom-3b")
model = BloomForCausalLM.from_pretrained("tomrb/bettercallbloom-3b",low_cpu_mem_usage=True)

generator = pipeline('text-generation', model=model, tokenizer=tokenizer,do_sample=False)


def preprocess(text):
    #We add 'Question :' and 'Answer #1:' at the start and end of the prompt
    return "\nQuestion: " + text + "\nAnswer #1:"


def generate(text):
    
    preprocessed_text = preprocess(text)
    result = generator(preprocessed_text, max_length=128)
    output = re.split(r'\nQuestion:|Answer #1:|Answer #|Title:',result[0]['generated_text'])[2]
    
    return output

examples = [
    ["I started a company with a friend. What types of legal documents should we fill in to clarify the ownership of the company?"],
    ["[CA] I got a parking ticket in Toronto. How can I contest it?"],
]

demo = gr.Interface(
    fn=generate,
    inputs=gr.inputs.Textbox(lines=5, label="Input Text", placeholder = "Write your question here..."),
    outputs=gr.outputs.Textbox(label="Generated Text"),
    examples=examples,
    description=description,
    title=title
)

demo.launch()