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?"],
]

with gr.Blocks() as demo:
  gr.Markdown("<h1><center>Better Call Bloom!</center></h1>")
  gr.Markdown("""<center>When in legal doubt, you better call BLOOM! Ask BLOOM any legal question:
<img src=https://huggingface.co/spaces/tomrb/bettercallbloom/resolve/main/img.jpeg width=200px></center>""")
  gr.Markdown("""<center>***THIS IS NOT LEGAL ADVICE. 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.***</center>""")
  

  input_text = gr.Textbox(label="Input", lines=6)  
  buton = gr.Button("Submit ")  
  output_text = gr.Textbox(lines=6, label="Output")
  buton.click(generate, inputs=[input_text], outputs=output_text)  

   
  gr.HTML("""
    <div style="border-top: 1px solid #303030;">
      <br>
      <p>Space by: <a href="https://twitter.com/TomRBeaudoin"><img src="https://img.shields.io/twitter/follow/TomRBeaudoin?label=%40thomasrbeaudoin&style=social" alt="Twitter Follow"></a></p><br>
      <p>Help me pay for GPU hours so I can publish faster models!</p>
      <a href="https://www.buymeacoffee.com/thomasrb" 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=tomrb.bettercallbloom" alt="visitors"></p>
    </div>
    """)
demo.launch(enable_queue=True, debug=True)