Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 model_name = "anasmkh/customized_llama3.1_8b"
+
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
 
@@ -11,20 +12,22 @@ generator = pipeline(
     tokenizer=tokenizer,
     max_new_tokens=64,
     temperature=1.5,
-    min_p=0.1
+    min_p=0.1,
 )
 
-def generate_response(prompt):
-  messages = [
-    {"role": "user", "content": prompt},
-  ]
-  response = generator(messages)[0]['generated_text']
-  return response.split("<|end_header_id|>")[1].strip()
+def chat(message, history):
+    history = history or []
+    history.append({"role": "user", "content": message})
+    response = generator(history)[-1]["generated_text"]
+    history.append({"role": "assistant", "content": response})
+    return history
 
-demo = gr.Interface(
-    fn=generate_response,
-    inputs=gr.Textbox(lines=5, label="Enter your prompt"),
-    outputs=gr.Textbox(label="Model Response")
-)
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot()
+    message = gr.Textbox()
+    clear = gr.ClearButton([message, chatbot])
+
+    message.submit(chat, [message, chatbot], chatbot)
+    clear.click(lambda: None, None, chatbot, queue=False)
 
 demo.launch()