Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 from transformers import pipeline
 from transformers import BloomTokenizerFast, BloomForCausalLM
-import re
 
 description = """
 When in legal doubt, you better call BLOOM! Ask BLOOM any legal question:
@@ -10,10 +9,12 @@ When in legal doubt, you better call BLOOM! Ask BLOOM any legal question:
 title = "Better Call Bloom!"
 examples = [["Adventurer is approached by a mysterious stranger in the tavern for a new quest."]]
 
-tokenizer = BloomTokenizerFast.from_pretrained("tomrb/bettercallbloom-3b")
-model_8bit = BloomForCausalLM.from_pretrained("tomrb/bettercallbloom-3b",device_map="auto",load_in_8bit=True)
 
-generator = pipeline('text-generation', model=model_8bit, tokenizer=tokenizer)
+
+tokenizer = BloomTokenizerFast.from_pretrained("tomrb/bettercallbloom-3b-8bit")
+model = BloomForCausalLM.from_pretrained("tomrb/bettercallbloom-3b-8bit")
+
+generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
 
 
 def preprocess(text):
@@ -24,9 +25,9 @@ def preprocess(text):
 def generate(text):
     
     preprocessed_text = preprocess(text)
-    result = generator(preprocessed_text, max_length=256)
-    #output = re.split(r'\nQuestion:|Answer #|Title:',result[0]['generated_text'])[2]
-    output = result[0]['generated_text']
+    result = generator(preprocessed_text, max_length=128)
+    output = re.split(r'\nQuestion:|Answer #|Title:',result[0]['generated_text'])[2]
+    
     return output
 
 examples = [