Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import os
 import requests
 import json
+import re
 
 #from huggingface_hub.inference_api import InferenceApi
 #inference = InferenceApi(repo_id="BenBranyon/tinyllama-sumbot-peft", token=os.environ['HF_TOKEN'])
@@ -17,13 +18,11 @@ def query(payload):
     response = requests.request("POST", API_URL, headers=headers, data=data)
     try:
         data = json.loads(response.content.decode("utf-8"))
-        response = data[0]['generated_text']
+        generated_text = data[0]['generated_text']
+        reponse = re.sub('\n<|system|>.*?<|assistant|>','',generated_text, flags=re.DOTALL)
         return response
     except:
         return response
-    #response = inference(payload)
-    #response = requests.post(API_URL, headers=headers, json=payload)
-    #return response.json()
 
 messages = st.container()
 message = st.chat_message("assistant").write("Hello human")