Update app.py

app.py

@@ -3,12 +3,10 @@ import requests
 import os
 import json
 import time
+from llama3 import llama3chat
 
 # Get the API key from the environment variable
 HF_TOKEN = os.getenv("HF_TOKEN")
-
-icons = {"assistant": "robot.png", "user": "man-kddi.png", "system": "man-kddi.png"}
-
 if HF_TOKEN is None:
     st.error("API key not found. Please set the HF_TOKEN secret in your Hugging Face Space.")
     st.stop()
@@ -16,69 +14,22 @@ if HF_TOKEN is None:
 # Hugging Face API settings
 API_URL = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-8B-Instruct"
 
-# Define the headers
-headers = {
-    'Content-Type': 'application/json',
-    'Authorization': f'Bearer {HF_TOKEN}'
-}
-
-
-def format_llama_messages(messages):
-    """Format messages for Llama-3 chat models.
-    """
-    prompt = []
-
-    if messages[0]["role"] == "system":
-        content = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n"+messages[0]["content"]+ "<|eot_id|>\n\n"
-        prompt.append(content)
-    for user, answer in zip(messages[1::2], messages[2::2]):
-        prompt.append("<|start_header_id|>user<|end_header_id|>"+ (user["content"]).strip()+ "<|eot_id|><|start_header_id|>assistant<|end_header_id|>"+ (answer["content"]).strip()+ "<|eot_id|>")
-
-    prompt.append("<|start_header_id|>user<|end_header_id|>"+ (messages[-1]["content"]).strip()+ "<|eot_id|>")
-
-    return "".join(prompt)
-
-def handle_query(query):
-    st.session_state.messages.append({'role': 'user', "content": query})
-    formatted_messages=format_llama_messages(st.session_state.messages)
-
-    json_body = {
-        "inputs": formatted_messages,
-                "parameters": {"max_new_tokens":max_new_tokens, "top_p":top_p, "temperature":temperature}
-        }
-    data = json.dumps(json_body)
-    print("JSON Sent:", data)
-    response = requests.request("POST", API_URL, headers=headers, data=data)
-    try:
-        data=json.loads(response.content.decode("utf-8"))
-        result=data[0]['generated_text']
-        last_comma_index = result.rfind('|eot_id|>assistant')
+icons = {"assistant": "robot.png", "user": "man-kddi.png", "system": "man-kddi.png"}
 
-        # If a comma is found, slice the string from the last comma onwards
-        if last_comma_index != -1:
-            result = result[last_comma_index + 18:].strip()
-            
-        return result
-    except:
-        return response
+_llama3chat = llama3chat(API_URL, HF_TOKEN, 1024, 0.9,0.5)
 
-def prompt_changed():
-    st.session_state.messages[0]["content"]=st.session_state['system_prompt']
-    
-def streamer(text):
-    for i in text:
-        yield i
-        time.sleep(0.005)
-        
 # Set up the Streamlit app
 st.title("Llama3 Chatbot")
 st.markdown("Using only request API calls, no external libraries here.")
 
 st.sidebar.markdown("Model Parameters")
-temperature = st.sidebar.slider("Temperature", 0.1, 1.0, 0.7, 0.1)
-max_new_tokens = st.sidebar.number_input("Max New Tokens", 50, 1024, 512, step=50)
-top_p = st.sidebar.slider("Top P", 0.1, 1.0, 0.9, 0.1)
+_llama3chat.temperature = st.sidebar.slider("Temperature", 0.1, 1.0, 0.7, 0.1)
+_llama3chat.max_new_tokens = st.sidebar.number_input("Max New Tokens", 50, 1024, 512, step=50)
+_llama3chat.top_p = st.sidebar.slider("Top P", 0.1, 1.0, 0.9, 0.1)
 
+def prompt_changed():
+    st.session_state.messages[0]["content"]=st.session_state['system_prompt']
+    
 system_prompt = st.sidebar.text_area(
     "System Prompt",
     "You are a food critic Chatbot. Reply in the style of a New York Times food critic",
@@ -94,15 +45,19 @@ for message in st.session_state.messages[1:]:
 
 user_prompt = st.chat_input("Start you conversation here...")
 
+def streamer(text):
+    for i in text:
+        yield i
+        time.sleep(0.005)
+
 if user_prompt:
+    st.session_state.messages.append({'role': 'user', "content": query})
     with st.chat_message("user", avatar=icons["user"]):
         st.write(user_prompt)
-
     # Trigger assistant's response retrieval and update UI
     with st.spinner("Thinking..."):
-        response = handle_query(user_prompt)
+        response = _llama3chat.query(st.session_state.messages)
     with st.chat_message("assistant", avatar=icons["assistant"]):
         st.write(streamer(response))
-
     st.session_state.messages.append({'role': 'assistant', "content": response})