Update app.py

app.py

@@ -1,26 +1,49 @@
+Import transformers and gradio
+import transformers
+import gradio as gr
+import git
+import os
+os.system("pip install --upgrade pip")
 
-
+#Load arabert preprocessor
+import git
+git.Git("arabert").clone("https://github.com/aub-mind/arabert")
 from arabert.preprocess import ArabertPreprocessor
+arabert_prep = ArabertPreprocessor(model_name="bert-base-arabert", keep_emojis=False)
+
+
+#Load Model
 from transformers import EncoderDecoderModel, AutoTokenizer
 tokenizer = AutoTokenizer.from_pretrained("tareknaous/bert2bert-empathetic-response-msa")
 model = EncoderDecoderModel.from_pretrained("tareknaous/bert2bert-empathetic-response-msa")
-model.to("cuda")
 model.eval()
-arabert_prep = ArabertPreprocessor(model_name="bert-base-arabert", keep_emojis=False)
-def generate_response(text):
+
+def generate_response(text, minimum_length, p, temperature):
   text_clean = arabert_prep.preprocess(text)
   inputs = tokenizer.encode_plus(text_clean,return_tensors='pt')
-  outputs = model.generate(input_ids = inputs.input_ids.to("cuda"),
-                   attention_mask = inputs.attention_mask.to("cuda"),
+  outputs = model.generate(input_ids = inputs.input_ids,
+                   attention_mask = inputs.attention_mask,
                    do_sample = True,
-                   min_length=10,
-                   top_k = 0,
-                   top_p = 0.9,
-                   temperature = 0.5)
+                   min_length=minimum_length,
+                   top_p = p,
+                   temperature = temperature)
   preds = tokenizer.batch_decode(outputs) 
   response = str(preds)
   response = response.replace("\'", '')
   response = response.replace("[[CLS]", '')
   response = response.replace("[SEP]]", '')
   response = str(arabert_prep.desegment(response))
-  return response
+  return response
+
+title = 'Empathetic Response Generation in Arabic'
+description = 'This demo is for a BERT2BERT model trained for single-turn open-domain empathetic dialogue response generation in Modern Standard Arabic'
+gr.Interface(fn=generate_response,
+              inputs=[
+          gr.inputs.Textbox(),
+          gr.inputs.Slider(5, 20, step=1, label='Minimum Output Length'),
+          gr.inputs.Slider(0.7, 1, step=0.1,  label='Top-P'),
+          gr.inputs.Slider(1, 3, step=0.1, label='Temperature'),
+          ],
+             outputs="text",
+             title=title,
+             description=description).launch()