Update app.py

app.py

@@ -5,50 +5,91 @@ from transformers import AutoTokenizer
 import re
 
 model_name = "Naseej/AskMe-Large"
-tokenizer = AutoTokenizer.from_pretrained(model_name , bos_token='<|startoftext|>',
-                                          eos_token='<|endoftext|>', pad_token='<|pad|>')
+tokenizer = AutoTokenizer.from_pretrained(model_name, bos_token='<|startoftext|>',
+                                         eos_token='<|endoftext|>', pad_token='<|pad|>')
 model = GPT2LMHeadModel.from_pretrained(model_name)
 model.resize_token_embeddings(len(tokenizer))
-generator = pipeline("text-generation",model=model,tokenizer=tokenizer)
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-def generate(text, num_beams=10, temperature=1, do_sample=True, top_k=40, top_p=60):
-    prompt = f'Prompt: {text}\nAnswer:'
+def generate_response(message, history, num_beams=4, temperature=0.99, do_sample=True, top_k=60, top_p=0.9):
+    prompt = f'Prompt: {message}\nAnswer:'
     pred_text = generator(prompt,
-      pad_token_id=tokenizer.eos_token_id,
-      num_beams=int(num_beams), 
-      max_length=1024,
-            min_length=0,           
-      temperature=temperature,
-      do_sample=do_sample,                    
-      top_p=top_p,
-      top_k=int(top_k),
-      repetition_penalty=3.0,
-      no_repeat_ngram_size=3)[0]['generated_text']
+                         pad_token_id=tokenizer.eos_token_id,
+                         num_beams=int(num_beams),
+                         max_length=1024,
+                         min_length=0,
+                         temperature=temperature,
+                         do_sample=do_sample,
+                         top_p=top_p,
+                         top_k=int(top_k),
+                         repetition_penalty=3.0,
+                         no_repeat_ngram_size=3)[0]['generated_text']
     try:
         pred_sentiment = re.findall("Answer:(.*)", pred_text, re.S)[-1]
     except:
-        pred_sentiment = "None"
+        pred_sentiment = "لم أستطع توليد إجابة. يرجى إعادة صياغة السؤال."
 
     return pred_sentiment
 
-examples = [
-    ["اكتب مقال عن الذكاء الصناعي"],
-    ["اكتب قصة قصيرة عن النجاح"],['كيف يمكن المحافظة على حياه صحية'],
-]
+# CSS for RTL support and styling
+css = """
+.gradio-container {direction: rtl;}
+.message.user {background-color: #2b5797; color: white; border-radius: 20px; padding: 8px 12px; margin-bottom: 8px; text-align: right;}
+.message.bot {background-color: #f0f0f0; color: black; border-radius: 20px; padding: 8px 12px; margin-bottom: 8px; text-align: right;}
+"""
 
-demo = gr.Interface(
-    css=".gradio-container {direction: rtl;}",
-    fn=generate,
-    inputs=[
-        gr.inputs.Textbox(lines=5, label="Input Text"),
-        gr.inputs.Number(default=4, label="num_beams"),
-        gr.inputs.Slider(minimum=0.1, maximum=2.0, default=0.99, step=0.01, label="temperature"),
-        gr.inputs.Checkbox(default = True , label="النمط الابداعي"),
-        gr.inputs.Number(default=60, label="top_k"),
-        gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.9, step=0.01, label="top_p"),
-    ],
-    outputs=gr.outputs.Textbox(label="Generated Text"),
-    examples=examples
-)
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("# نظام AskMe - تحدث معي")
+    
+    with gr.Row():
+        with gr.Column(scale=3):
+            chatbot = gr.Chatbot(label="المحادثة", elem_classes=["chatbot"])
+            msg = gr.Textbox(label="اكتب رسالتك هنا", placeholder="اكتب هنا...")
+            
+            with gr.Row():
+                submit_btn = gr.Button("إرسال", variant="primary")
+                clear_btn = gr.Button("مسح المحادثة")
+        
+        with gr.Column(scale=1):
+            with gr.Accordion("إعدادات توليد النص", open=False):
+                num_beams = gr.Slider(1, 10, value=4, step=1, label="عدد الشعاعات")
+                temperature = gr.Slider(0.1, 2.0, value=0.99, step=0.01, label="درجة الحرارة")
+                do_sample = gr.Checkbox(value=True, label="النمط الإبداعي")
+                top_k = gr.Slider(1, 100, value=60, step=1, label="Top-K")
+                top_p = gr.Slider(0.0, 1.0, value=0.9, step=0.01, label="Top-P")
+    
+    # Example prompts
+    with gr.Accordion("أمثلة", open=False):
+        gr.Examples(
+            [
+                ["اكتب مقال عن الذكاء الصناعي"],
+                ["اكتب قصة قصيرة عن النجاح"],
+                ["كيف يمكن المحافظة على حياه صحية"]
+            ],
+            inputs=msg
+        )
+    
+    # Set up event handlers
+    submit_btn.click(
+        generate_response,
+        inputs=[msg, chatbot, num_beams, temperature, do_sample, top_k, top_p],
+        outputs=[chatbot],
+    ).then(
+        lambda: "", 
+        None, 
+        [msg]
+    )
+    
+    msg.submit(
+        generate_response,
+        inputs=[msg, chatbot, num_beams, temperature, do_sample, top_k, top_p],
+        outputs=[chatbot],
+    ).then(
+        lambda: "", 
+        None, 
+        [msg]
+    )
+    
+    clear_btn.click(lambda: None, None, chatbot, queue=False)
 
-demo.launch()
+demo.launch()