improve the UI

app.py

@@ -1,6 +1,7 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import gradio as gr
+import re
 
 # Check if GPU is available
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -8,7 +9,6 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # Load the model and tokenizer with GPU optimizations
 model = AutoModelForCausalLM.from_pretrained(
     "AlanYky/phi-3.5_tweets_instruct",
-    device_map="cpu",
     torch_dtype=torch.float16,  # Use FP16 for GPU
     trust_remote_code=True,
     load_in_8bit=True
@@ -37,17 +37,53 @@ generation_args = {
 
 
 # Function for tweet generation
-def generate_tweet(instruction):
-    messages = [{"role": "user", "content": instruction}]
+def generate_tweet(tweet_idea):
+    max_length = 50  # Set a maximum input length
+    if len(tweet_idea) > max_length:
+        return f"Error: Tweet idea exceeds the maximum allowed length of {max_length} characters."
+
+    prompt = f"Could you generate a tweet about {tweet_idea}?"
+    messages = [{"role": "user", "content": prompt}]
+    output = pipe(messages, **generation_args)
+    return clean_tweet(output[0]['generated_text'])
+
+
+def generate_tweet_with_rag(tweet_idea):
+    prompt = f"Could you generate a tweet about {tweet_idea}?"
+    messages = [{"role": "user", "content": prompt}]
     output = pipe(messages, **generation_args)
     return output[0]['generated_text']
 
 
+def clean_tweet(tweet):
+    # Remove trailing sequences of non-ASCII characters (like `����������������`)
+    cleaned_tweet = tweet.replace("�", "")
+
+    cleaned_tweet = re.sub(r"\(\)", "", cleaned_tweet)
+
+    # change more than one dot to one dot
+    cleaned_tweet = re.sub(r"\.{4,}", "...", cleaned_tweet)
+
+    # change more than two question marks to two question marks
+    cleaned_tweet = re.sub(r"\?{4,}", "???", cleaned_tweet)
+
+    # change more than two ! to two !
+    cleaned_tweet = re.sub(r"!{4,}", "!!!", cleaned_tweet)
+
+    # Remove repeated double quotes at the end
+    cleaned_tweet = re.sub(r'"+$', '', cleaned_tweet)
+
+    return cleaned_tweet.strip()
+
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Tweet Generator (Optimized for GPU)")
+    gr.Markdown("# Tweet Generator")
     with gr.Row():
-        instruction_input = gr.Textbox(label="Instruction", placeholder="Generate a tweet about...")
+        instruction_input = gr.Textbox(
+            label="Instruction",
+            placeholder="Enter your tweet idea (It can be a topic, hashtag, sentence, or any format)..."
+        )
+
     generate_button = gr.Button("Generate")
     output_box = gr.Textbox(label="Generated Tweet", placeholder="Your tweet will appear here.")