Update chatbot_bedrock.py

add Hugging Face runnable model wrapper

chatbot_bedrock.py

@@ -6,32 +6,28 @@ import langchain.globals
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import streamlit as st
 
-@st.cache_resource 
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("KvrParaskevi/Hotel-Assistant-Attempt4-Llama-2-7b")
-    model = AutoModelForCausalLM.from_pretrained("KvrParaskevi/Hotel-Assistant-Attempt4-Llama-2-7b")
-    return tokenizer,model
-
+class HuggingFaceModelWrapper(Runnable):  # Assuming Runnable is the required interface
+    def __init__(self, model, tokenizer):
+        self.model = model
+        self.tokenizer = tokenizer
 
-#Write function to connect to Bedrock
-# def demo_chatbot():
-#     # client = boto3.client('bedrock-runtime')
-    
-#     template = """Question: {question}
+    def run(self, input_text):
+        # Convert the input text to tokens
+        input_ids = self.tokenizer.encode(input_text, return_tensors="pt")
 
-#     Answer: Let's think step by step."""
-#     prompt = PromptTemplate(template=template, input_variables=["question"])
-#     llm=HuggingFaceHub(repo_id="google/flan-t5-xl", model_kwargs={"temperature":1e-10})
+        # Generate a response from the model
+        output = self.model.generate(input_ids, max_length=100, num_return_sequences=1)
 
-#     question = "When was Google founded?"
+        # Decode the generated tokens to a string
+        response_text = self.tokenizer.decode(output[0], skip_special_tokens=True)
+        return response_text
 
-#     print(llm_chain.run(question))
-#     return demo_llm
 
-#test out the code with the Predicgt method
-    #return demo_llm.predict(input)
-# = demo_chatbot('What is the temperature in Nuremberg today?')
-#print(response)
+@st.cache_resource 
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained("KvrParaskevi/Hotel-Assistant-Attempt4-Llama-2-7b")
+    model = AutoModelForCausalLM.from_pretrained("KvrParaskevi/Hotel-Assistant-Attempt4-Llama-2-7b")
+    return tokenizer,model
 
 def demo_miny_memory(model):
     # llm_data = get_Model(hugging_face_key)