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 legal 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
)
gr.HTML("""
    <div style="border-top: 1px solid #303030;">
      <br>
      <p>Space by: <a href="https://twitter.com/hahahahohohe"><img src="https://img.shields.io/twitter/follow/hahahahohohe?label=%40anzorq&style=social" alt="Twitter Follow"></a></p><br>
      <a href="https://www.buymeacoffee.com/anzorq" target="_blank"><img src="https://cdn.buymeacoffee.com/buttons/v2/default-yellow.png" alt="Buy Me A Coffee" style="height: 45px !important;width: 162px !important;" ></a><br><br>
      <p><img src="https://visitor-badge.glitch.me/badge?page_id=anzorq.sd-to-diffusers" alt="visitors"></p>
    </div>
    """)

demo.launch()