Update app.py

app.py

@@ -1,59 +1,80 @@
-import gradio as gr
-import spaces
-from langchain_core.pydantic_v1 import BaseModel, Field
-from langchain.prompts import HumanMessagePromptTemplate, ChatPromptTemplate, SystemMessagePromptTemplate
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from langchain import HuggingFaceHub
+from langchain.llms.base import LLM
+from langchain.memory import ConversationBufferMemory,ConversationBufferWindowMemory
+from langchain.chains import LLMChain, ConversationChain
+from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
+from langchain_community.llms import HuggingFaceEndpoint
+from langchain.prompts import PromptTemplate, ChatPromptTemplate
+from langchain_core.prompts.chat import SystemMessagePromptTemplate, HumanMessagePromptTemplate
+import streamlit as st
 
-# Load the Hugging Face model and tokenizer
-model_name = "KvrParaskevi/Llama-2-7b-Hotel-Booking-Model"
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+your_endpoint_url = "https://kp4xdy196cw81uf3.us-east-1.aws.endpoints.huggingface.cloud"
+token = st.secrets["HUGGINGFACEHUB_API_TOKEN"]
 
-system_message = """<<SYS>>
-You are an AI having conversation with a human. Below is an instruction that describes a task. 
-Write a response that appropriately completes the request.
-Reply with the most helpful and logic answer. During the conversation you need to ask the user 
-the following questions to complete the hotel booking task.
-1) Where would you like to stay and when?
-2) How many people are staying in the room?
-3) Do you prefer any ammenities like breakfast included or gym?
-4) What is your name, your email address and phone number? 
-Make sure you receive a logical answer from the user from every question to complete the hotel 
-booking process.
-<</SYS>>
-"""
-messages = [
-    SystemMessagePromptTemplate.from_template(system_message)
-]
+llm = HuggingFaceEndpoint(
+    endpoint_url=f"{your_endpoint_url}",
+    huggingfacehub_api_token = f"{token}",
+    task = "text-generation",
+    max_new_tokens=128,
+    top_k=10,
+    top_p=0.95,
+    typical_p=0.95,
+    temperature=0.01,
+    repetition_penalty=1.03
+)
+#print(llm)
 
-# Define the Langchain chatbot function
-@spaces.GPU
-def chatbot(message, history):
-    # Create a Langchain prompt template
-    prompt_template = HumanMessagePromptTemplate.from_template(message)
-    # Create a Langchain chat prompt template
-    messages.append(prompt_template)
-    chat_prompt_template = ChatPromptTemplate.from_messages(messages)
-    # Use the Langchain TextIteratorStreamer to generate responses
-    streamer = TextStreamer(model, tokenizer)
-    response = model.generate(chat_prompt_template, streamer=streamer, max_new_tokens=20)
-    return response
+def chat_template_prompt():
+    template = """
+    Do not repeat questions and do not generate answer for user/human.
+    
+    You are a helpful hotel booking asssitant.
+    Below is an instruction that describes a task.
+    Write a response that appropriately completes the request.
+    Reply with the most helpful and logic answer. During the conversation you need to ask the user
+    the following questions to complete the hotel booking task.
+    1) Where would you like to stay and when?
+    2) How many people are staying in the room?
+    3) Do you prefer any ammenities like breakfast included or gym?
+    4) What is your name, your email address and phone number?
+    
+    When the booking task is completed, respond with "Thank you for choosing us.".
+    
+    {history}
+    
+    """
+    
+    system_prompt = SystemMessagePromptTemplate.from_template(template)
+    human_prompt = HumanMessagePromptTemplate.from_template("{input}")
+    chat_prompt = ChatPromptTemplate.from_messages([system_prompt, human_prompt])
+    return chat_prompt
+
+def chain():
+    #memory = ConversationBufferMemory(memory_key="history")
+    chat_prompt = chat_template_prompt()
+    memory = ConversationBufferWindowMemory(k=3) #memory_key="history"
+    llm_chain = LLMChain(llm=llm, memory = memory, prompt = chat_prompt)
+    memory.load_memory_variables({}) #Initialize memory
+    return llm_chain
+
+def chat_output(inputs):
+    llm_chaim = chain()
+    result = llm_chaim.predict(input = inputs)
+    return result
 
-# Create a Gradio chatbot interface
 with gr.Blocks() as demo:
-    #chatbot_interface = gr.Chatbot()
-    #msg = gr.Textbox()
-    #clear = gr.Button("Clear")
+    
+    chatbot_component = gr.Chatbot(height=300, label = "history")
+    textbox_component = gr.Textbox(placeholder="Can I help you to book a hotel?", container=False, label = "input", scale=7)
 
-    # Define the chatbot function as a Gradio interface
     demo.chatbot_interface = gr.Interface(
-        fn=chatbot,
+        fn=chat_output, 
         inputs=[
-            gr.Textbox(lines=1, label="Input"),
+            textbox_component
         ], 
-        outputs="text",
-        title="Langchain Chatbot",
-        description="A simple chatbot using Langchain and Hugging Face"
+        outputs=chatbot_component,
+        title = "Hotel Booking Assistant Chat 🤗",
+        description = "I am your hotel booking assistant. Feel free to start chatting with me."
     )
-
 demo.launch()