Update app.py

app.py

@@ -25,8 +25,9 @@ def generate_text(prompt, length=100, theme=None, **kwargs):
 
     # If a theme is specified, add it to the prompt as a prefix for a special token
     if theme:
-        prompt = ' <{}> '.format(theme.strip()) + prompt.strip()
+        prompt = f"<{theme.strip()}> {prompt.strip()}"
 
+    # Encode the input prompt
     input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
     attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=device)
     pad_token_id = tokenizer.eos_token_id
@@ -34,6 +35,7 @@ def generate_text(prompt, length=100, theme=None, **kwargs):
     # Set the max length of the generated text based on the input parameter
     max_length = length if length > 0 else 100
 
+    # Generate the text using the model
     sample_outputs = model.generate(
         input_ids,
         attention_mask=attention_mask,
@@ -47,38 +49,28 @@ def generate_text(prompt, length=100, theme=None, **kwargs):
         no_repeat_ngram_size=2,
         repetition_penalty=1.5,
     )
-    generated_text = tokenizer.decode(sample_outputs[0], skip_special_tokens=True)
-
-    # Post preprocessing of the generated text
-    
-    # Remove any leading and trailing quotation marks
-    generated_text = generated_text.strip('"')
-
-    # Remove leading and trailing whitespace
-    generated_text = generated_text.strip()
 
-    # Find the special token in the generated text and remove it
-    match = re.search(r'<([^>]+)>', generated_text)
-    if match:
-        generated_text = generated_text[:match.start()] + generated_text[match.end():]
-
-    # Remove any leading numeric characters and quotation marks
-    generated_text = re.sub(r'^\d+', '', generated_text)
-    generated_text = re.sub(r'^"', '', generated_text)
+    # Decode the generated text
+    generated_text = tokenizer.decode(sample_outputs[0], skip_special_tokens=True)
 
-    # Remove any newline characters from the generated text
-    generated_text = generated_text.replace('\n', '')
+    # Postprocessing of the generated text
+    generated_text = generated_text.strip().strip('"') # Remove leading and trailing whitespace, remove any leading and trailing quotation marks
+    generated_text = re.sub(r'<([^>]+)>', '', generated_text) # Find the special token in the generated text and remove it
+    generated_text = re.sub(r'^\d+|^"', '', generated_text) # Remove any leading numeric characters and quotation marks
+    generated_text = generated_text.replace('\n', '') # Remove any newline characters from the generated text
+    generated_text = re.sub(r'[^\w\s]+', '', generated_text) # Remove any other unwanted special characters
+    generated_text = generated_text.capitalize()
 
-    # Remove any other unwanted special characters
-    generated_text = re.sub(r'[^\w\s]+', '', generated_text)
-    
-    return generated_text.strip().capitalize()
+    return generated_text
 
-# Define a Gradio interface for the generate_text function, allowing users to input a prompt and generate text based on it
+# Define a Gradio interface for the generate_text function
 iface = gr.Interface(
     fn=generate_text,
-    inputs=['text', gr.inputs.Slider(minimum=10, maximum=100, default=50, label='Length of text'),
-            gr.inputs.Textbox(default='Food', label='Theme')],
+    inputs=[
+        "text", 
+        gr.inputs.Slider(minimum=10, maximum=100, default=50, label='Length of text'),
+        gr.inputs.Textbox(default='Food', label='Theme')
+    ],
     outputs=[gr.outputs.Textbox(label='Generated Text')],
     title='Yelp Review Generator',
     description='Generate a Yelp review based on a prompt, length of text, and theme.',
@@ -94,4 +86,4 @@ iface = gr.Interface(
     flagging_options=[("🙌", "positive"), ("😞", "negative")],
 )
 
-iface.launch(debug=False)
+iface.launch(debug=False, share=True)