Update app.py

app.py

@@ -14,18 +14,29 @@ def initialize_model_and_tokenizer(model_name="KvrParaskevi/Llama-2-7b-Hotel-Boo
 
 model, tokenizer = initialize_model_and_tokenizer()
 
-class CustomLLM(LLM):
-    def _call(self, prompt, stop=None, run_manager=None) -> str:
-        inputs = tokenizer(prompt, return_tensors="pt")
-        result = model.generate(input_ids=inputs.input_ids, max_new_tokens=20)
-        result = tokenizer.decode(result[0])
-        return result
+def load_pipeline():
+    tokenizer, model = load_model()
+    pipe = pipeline("text-generation", 
+                    model= model, 
+                    tokenizer = tokenizer,
+                    max_new_tokens = 20,
+                    top_k = 30, 
+                    early_stopping=True,
+                    num_beams = 2,
+                    temperature = 0.1,
+                    repetition_penalty = 1.03)
 
-    @property
-    def _llm_type(self) -> str:
-        return "custom"
+    llm = HuggingFacePipeline(pipeline = pipe)
+    return llm
 
-llm = CustomLLM()
+def chat_interface(inputs):
+    question = inputs["input"]
+    chat_history = inputs["history"]
+    # Assuming `chain` is your instance of ConversationalRetrievalChain
+    result = chain.run({"input": question, "history": chat_history})
+    return result["response"]
+
+llm = load_pipeline()
 
 template = """<<SYS>>
 You are an AI having conversation with a human. Below is an instruction that describes a task. 
@@ -49,24 +60,27 @@ memory = ConversationBufferMemory(memory_key="history", llm = llm, prompt = prom
 llm_chain = ConversationChain(prompt=prompt, llm=llm, memory = memory)
 
 with gr.Blocks() as demo:
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox()
+    gr.Markdown("Hotel Booking Assistant Chat 🤗")
+    chatbot = gr.Chatbot(label="Chat history")
+    message = gr.Textbox(label="Ask me a question!")
     clear = gr.Button("Clear")
     #llm_chain, llm = init_chain(model, tokenizer)
 
-    def user(user_message, history):
-        return "", history + [[user_message, None]]
-    @spaces.GPU
-    def bot(history):
-        print("Question: ", history[-1][0])
-        bot_message = llm_chain.invoke(input=history[-1][0])
-        print("Response: ", bot_message)
-        history[-1][1] = ""
-        history[-1][1] += bot_message
-        return history
+    # Convert chat history to list of tuples
+    chat_history_tuples = []
+    for message in chat_history:
+        chat_history_tuples.append((message[0], message[1]))
+    
+    result = llm_chain({"input": query, "history": chat_history_tuples})
+    
+    iface = gr.Interface(
+        fn=chat_interface, 
+        inputs=[
+            gr.inputs.Textbox(lines=1, label="Question"),
+            gr.inputs.Textbox(lines=5, label="Chat History"),
+        ], 
+        outputs="text"
+    )
+iface.launch()
 
-    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()