updates

app.py

@@ -26,9 +26,10 @@ DESCRIPTION = """
 """
 
 EXAMPLES = [
-    ["What is the capital of France?", "You are a helpful assistant.", []],
-    ["Explain the theory of relativity in simple terms.", "You are an expert physicist explaining concepts to a layman.", []],
-    ["Write a short poem about artificial intelligence.", "You are a creative poet with a interest in technology.", []]
+    ["What is the capital of France?", "You are a helpful assistant."],
+    ["Who landed on the moon?", "You are a knowledgeable historian."],
+    ["Who is Yarin Gal?", "You are a helpful assistant."],
+    ["Explain the theory of relativity in simple terms.", "You are an expert physicist explaining concepts to a layman."],
 ]
 
 if torch.cuda.is_available():
@@ -90,22 +91,7 @@ class CustomStreamer(TextIteratorStreamer):
 #         gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
 #     input_ids = input_ids.to(model.device)
 
-#     streamer = CustomStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-#     generation_kwargs = dict(
-#         input_ids=input_ids,
-#         max_new_tokens=max_new_tokens,
-#         do_sample=True,
-#         top_p=top_p,
-#         top_k=top_k,
-#         temperature=temperature,
-#         repetition_penalty=repetition_penalty,
-#         streamer=streamer,
-#         output_hidden_states=True,
-#         return_dict_in_generate=True,
-#     )
-
-#     thread = Thread(target=model.generate, kwargs=generation_kwargs)
-#     thread.start()
+
 
 #     se_highlighted_text = ""
 #     acc_highlighted_text = ""
@@ -133,7 +119,6 @@ class CustomStreamer(TextIteratorStreamer):
 @spaces.GPU
 def generate(
     message: str,
-    chat_history: List[Tuple[str, str]],
     system_prompt: str,
     max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
     temperature: float = 0.6,
@@ -144,8 +129,6 @@ def generate(
     conversation = []
     if system_prompt:
         conversation.append({"role": "system", "content": system_prompt})
-    for user, assistant in chat_history:
-        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
     conversation.append({"role": "user", "content": message})
 
     input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
@@ -200,7 +183,9 @@ def generate(
         se_highlighted_text += f" {se_new_highlighted_text}"
         acc_highlighted_text += f" {acc_new_highlighted_text}"
         
-        yield se_highlighted_text, acc_highlighted_text
+        # yield se_highlighted_text, acc_highlighted_text
+    return se_highlighted_text, acc_highlighted_text
+    
 
 
 def highlight_text(text: str, uncertainty_score: float) -> str:
@@ -237,27 +222,33 @@ with gr.Blocks(title="Llama-2 7B Chat with Dual Probes", css="footer {visibility
     
     with gr.Row():
         generate_btn = gr.Button("Generate")
-    # add spacing between probes and titles for each output
+    # Add spacing between probes
+    gr.HTML("<br><br>")
+
     with gr.Row():
         with gr.Column():
+            # make a box
             title = gr.HTML("<h2>Semantic Uncertainty Probe</h2>")
             se_output = gr.HTML(label="Semantic Uncertainty Probe")
+        
+        # Add spacing between columns
+        gr.HTML("<div style='width: 20px;'></div>")
+
         with gr.Column():
             title = gr.HTML("<h2>Accuracy Probe</h2>")
             acc_output = gr.HTML(label="Accuracy Probe")
     
-    chat_history = gr.State([])
-    
     gr.Examples(
         examples=EXAMPLES,
-        inputs=[message, system_prompt, chat_history],
+        inputs=[message, system_prompt],
         outputs=[se_output, acc_output],
         fn=generate,
+
     )
     
     generate_btn.click(
         generate,
-        inputs=[message, system_prompt, chat_history, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
+        inputs=[message, system_prompt, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
         outputs=[se_output, acc_output]
     )