Update app.py

app.py

@@ -5,6 +5,8 @@ import gradio as gr
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from typing import List, Dict, Optional, Tuple
+from http import HTTPStatus
 
 DESCRIPTION = """
 # QwQ Distill
@@ -44,21 +46,60 @@ model.eval()
 if tokenizer.pad_token_id is None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
+# Define roles for the chat
+class Role:
+    SYSTEM = "system"
+    USER = "user"
+    ASSISTANT = "assistant"
+
+# Default system message
+default_system = "You are a helpful assistant."
+
+def clear_session() -> List:
+    return "", []
+
+def modify_system_session(system: str) -> Tuple[str, str, List]:
+    if system is None or len(system) == 0:
+        system = default_system
+    return system, system, []
+
+def history_to_messages(history: List, system: str) -> List[Dict]:
+    messages = [{'role': Role.SYSTEM, 'content': system}]
+    for h in history:
+        messages.append({'role': Role.USER, 'content': h[0]})
+        messages.append({'role': Role.ASSISTANT, 'content': h[1]})
+    return messages
+
+def messages_to_history(messages: List[Dict]) -> Tuple[str, List]:
+    assert messages[0]['role'] == Role.SYSTEM
+    system = messages[0]['content']
+    history = []
+    for q, r in zip(messages[1::2], messages[2::2]):
+        history.append([q['content'], r['content']])
+    return system, history
+
 @spaces.GPU(duration=120)
 def generate(
-    message: str,
-    chat_history: list[dict],
+    query: Optional[str],
+    history: Optional[List],
+    system: str,
     max_new_tokens: int = 1024,
     temperature: float = 0.6,
     top_p: float = 0.9,
     top_k: int = 50,
     repetition_penalty: float = 1.2,
-) -> Iterator[str]:
-    conversation = chat_history.copy()
-    conversation.append({"role": "user", "content": message})
+) -> Iterator[Tuple[str, List, str]]:
+    if query is None:
+        query = ''
+    if history is None:
+        history = []
+
+    # Convert history to messages
+    messages = history_to_messages(history, system)
+    messages.append({'role': Role.USER, 'content': query})
 
     # Apply chat template and get input_ids
-    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
+    input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
 
     # Create attention mask
     attention_mask = torch.ones_like(input_ids)
@@ -94,12 +135,17 @@ def generate(
     outputs = []
     for text in streamer:
         outputs.append(text)
-        yield "".join(outputs)
+        response = "".join(outputs)
+        # Update history with the new response
+        new_messages = messages + [{'role': Role.ASSISTANT, 'content': response}]
+        system, new_history = messages_to_history(new_messages)
+        yield "", new_history, system
 
 
 demo = gr.ChatInterface(
     fn=generate,
     additional_inputs=[
+        gr.Textbox(label="System Message", value=default_system, lines=2),
         gr.Slider(
             label="Max new tokens",
             minimum=1,
@@ -138,13 +184,12 @@ demo = gr.ChatInterface(
     ],
     stop_btn=None,
     examples=[
-        ["Write a Python function to reverses a string if it's length is a multiple of 4. def reverse_string(str1): if len(str1) % 4 == 0: return ''.join(reversed(str1)) return str1 print(reverse_string('abcd')) print(reverse_string('python')) "],
-        ["Rectangle $ABCD$ is the base of pyramid $PABCD$. If $AB = 10$, $BC = 5$, $\overline{PA}\perp \text{plane } ABCD$, and $PA = 8$, then what is the volume of $PABCD$?"],
-        ["Difference between List comprehension and Lambda in Python lst  =  [x ** 2  for x in range (1, 11)   if  x % 2 == 1] print(lst)"],
+        ["Write a Python function to reverses a string if it's length is a multiple of 4."],
+        ["What is the volume of a pyramid with a rectangular base?"],
+        ["Explain the difference between List comprehension and Lambda in Python."],
         ["What happens when the sun goes down?"],
     ],
     cache_examples=False,
-    type="messages",
     description=DESCRIPTION,
     css=css,
     fill_height=True,
@@ -152,4 +197,4 @@ demo = gr.ChatInterface(
 
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=20).launch(share=True)