Update app.py

app.py

@@ -1,4 +1,33 @@
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from langchain.llms.base import LLM
+from langchain import PromptTemplate, LLMChain
+
+def initialize_model_and_tokenizer(model_name="KvrParaskevi/Llama-2-7b-Hotel-Booking-Model"):
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    return model, tokenizer
+
+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
+
+    @property
+    def _llm_type(self) -> str:
+        return "custom"
+
+llm = CustomLLM()
+
+template = """Question: {question}
+Answer: Let's think step by step."""
+prompt = PromptTemplate(template=template, input_variables=["question"])
+
+llm_chain = LLMChain(prompt=prompt, llm=llm)
 
 with gr.Blocks() as demo:
     chatbot = gr.Chatbot()