Update chatbot_bedrock.py

chatbot_bedrock.py

@@ -5,37 +5,25 @@ from langchain.chains import ConversationChain
 import langchain.globals
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import streamlit as st
-from langchain_core.runnables.base import Runnable
+from langchain_community.llms import HuggingFaceHub
 
-class HuggingFaceModelWrapper(Runnable):  # Assuming Runnable is the required interface
-    def __init__(self, model, tokenizer):
-        self.model = model
-        self.tokenizer = tokenizer
-
-    def run(self, input_text):
-        # Convert the input text to tokens
-        input_ids = self.tokenizer.encode(input_text, return_tensors="pt")
-
-        # Generate a response from the model
-        output = self.model.generate(input_ids, max_length=100, num_return_sequences=1)
-
-        # Decode the generated tokens to a string
-        response_text = self.tokenizer.decode(output[0], skip_special_tokens=True)
-        return response_text
-    def invoke(self, *args, **kwargs):
-        # Implement the 'invoke' method as required by the abstract base class/interface
-        # The implementation here depends on what 'invoke' is supposed to do. As an example:
-    
-        # Assuming 'invoke' should process some input and return a model response
-        input_text = args[0] if args else kwargs.get('input_text', '')
-        return self.run(input_text)
 
 
 @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
+    #tokenizer = AutoTokenizer.from_pretrained("KvrParaskevi/Hotel-Assistant-Attempt4-Llama-2-7b")
+    #model = AutoModelForCausalLM.from_pretrained("KvrParaskevi/Hotel-Assistant-Attempt4-Llama-2-7b")
+    model = HuggingFaceHub(
+        repo_id="KvrParaskevi/Hotel-Assistant-Attempt4-Llama-2-7b",
+        task="text-generation",
+        model_kwargs={
+            "max_new_tokens": 512,
+            "top_k": 30,
+            "temperature": 0.1,
+            "repetition_penalty": 1.03,
+        },
+    )
+    return model
 
 def demo_miny_memory(model):
     # llm_data = get_Model(hugging_face_key)