updated model and interface

app.py

@@ -48,8 +48,8 @@ tokenizer.pad_token = tokenizer.eos_token
 model = PeftModel.from_pretrained(model, PEFT_MODEL)
 
 generation_config = model.generation_config
-generation_config.max_new_tokens = 200
-generation_config.temperature = 0.7
+generation_config.max_new_tokens = 150
+generation_config.temperature = 0.6
 generation_config.top_p = 0.7
 generation_config.num_return_sequences = 1
 generation_config.pad_token_id = tokenizer.eos_token_id
@@ -61,21 +61,58 @@ pipeline = transformers.pipeline(
     tokenizer=tokenizer,
 )
 
-def query_model(message, history):
+def main():
 
-    prompt = f"""
-    <human>: {message}
-    <assistant>:
-    """.strip()
-    
-    result = pipeline(
-        prompt,
-        generation_config=generation_config,
-    )
+    with gr.Blocks() as demo:
 
-    # parsed_result = result[0]["generated_text"].split("<assistant>:")[1][1:]
+        def update_temp(temp):
+            generation_config.temperature = temp
 
-    return parsed_result
+        def update_tokens(tokens):
+            generation_config.max_new_tokens = tokens
+        
+        chatbot = gr.Chatbot(label="Fetch Rewards Chatbot")
+        temperature = gr.Slider(0, 1, value=0.6, step=0.1, label="Creativity", interactive=True)
+        temperature.change(fn=update_temp, inputs=temperature)
 
-    
-gr.ChatInterface(query_model, textbox=gr.Textbox(placeholder="Ask anything about Fetch!", container=False, scale=7),).launch()
+        tokens = gr.Slider(50, 200, value=100, step=50, label="Length", interactive=True)
+        tokens.change(fn=update_tokens, inputs=tokens)
+
+        msg = gr.Textbox(label="", placeholder="Ask anything about Fetch!")
+        clear = gr.Button("Clear Log")
+
+        def user(user_message, history):
+            return "", history + [[user_message, None]]
+
+        def bot(history):
+            
+            message = history[-1][0]
+            prompt = f"""
+            <human>: {message}
+            <assistant>:
+            """.strip()
+            
+            result = pipeline(
+                prompt,
+                generation_config=generation_config,
+            )
+            # print(result)
+            parsed_result = result[0]["generated_text"].split("<assistant>:")[1][1:].split("\n")[0]
+                
+            history[-1][1] = ""
+            for character in parsed_result:
+                history[-1][1] += character
+                time.sleep(0.01)
+                yield history
+
+        msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+            bot, chatbot, chatbot
+        )
+        clear.click(lambda: None, None, chatbot, queue=False)
+
+    demo.queue()
+    demo.launch()
+
+
+if __name__ == "__main__":
+    main()