Update app.py

app.py

@@ -3,30 +3,51 @@ import gradio as gr
 def generate_ielts_essay(prompt):
     # You may want to customize the max_length and other parameters according to your needs
     return "Hello " + prompt + "!!"
-    
-def generate_or_grade(task, text):
-    if task == "Generate":
-        # Generate an essay
-        # output = generator(text, max_length=500)[0]['generated_text']
-        return "Hello " + text + "!!"
-    elif task == "Grade":
-        # Simple grading function. Replace with a more sophisticated essay grading algorithm.
-        words = text.split()
-        if len(words) > 200:
-            return "Good"
-        else:
-            return "Poor"
-
-iface = gr.Interface(
-    fn=generate_or_grade, 
-    inputs=[
-        gr.inputs.Radio(["Generate", "Grade"], label="Task"), 
-        gr.inputs.Textbox(lines=10, placeholder="Enter text...")
-    ],
-    outputs="text",
-    title="IELTS Essay Generator and Grader",
-    description="This tool generates a response to an IELTS writing prompt using AI and provides a simple grading.",
-)
-
-iface.launch()
+    def format_chat_prompt(message, band, chat_history, instruction):
+    prompt = f"System:{instruction}"
+    for turn in chat_history:
+        user_message, bot_message = turn
+        prompt = f"{prompt}\nUser: {user_message}\nAssistant: {bot_message}"
+    prompt = f"{prompt}\nUser: IELTS Writing Band {band} - {message}\nAssistant:"
+    return prompt
+
+def respond(message, band, chat_history, instruction, temperature=0.7):
+    prompt = format_chat_prompt(message, band, chat_history, instruction)
+    chat_history = chat_history + [[message, ""]]
+    max_tokens = {6: 100, 7: 200, 8: 300, 9: 400}[int(band)]  # Change these values as needed
+    stream = client.generate_stream(prompt,
+                                    max_new_tokens=max_tokens,
+                                    stop_sequences=["\nUser:", ""],
+                                    temperature=temperature)
+    acc_text = ""
+    for idx, response in enumerate(stream):
+            text_token = response.token.text
+
+            if response.details:
+                return
+
+            if idx == 0 and text_token.startswith(" "):
+                text_token = text_token[1:]
+
+            acc_text += text_token
+            last_turn = list(chat_history.pop(-1))
+            last_turn[-1] += acc_text
+            chat_history = chat_history + [last_turn]
+            yield "", chat_history
+            acc_text = ""
+
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot(height=240)
+    msg = gr.Textbox(label="IELTS Writing Prompt")
+    band = gr.Radio([6, 7, 8, 9], label="Band")
+    with gr.Accordion(label="Advanced options",open=False):
+        system = gr.Textbox(label="System message", lines=2, value="A conversation between a user and an LLM-based AI assistant. The assistant writes IELTS articles of various bands.")
+        temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1, value=0.7, step=0.1)
+    btn = gr.Button("Submit")
+    clear = gr.ClearButton(components=[msg, chatbot], value="Clear console")
+
+    btn.click(respond, inputs=[msg, band, chatbot, system], outputs=[msg, chatbot])
+    msg.submit(respond, inputs=[msg, band, chatbot, system], outputs=[msg, chatbot]) 
+gr.close_all()
+demo.queue().launch(share=True)