Update app.py

app.py

@@ -11,13 +11,12 @@ from langchain_community.embeddings import HuggingFaceBgeEmbeddings
 from langchain_community.llms import CTransformers
 from ctransformers import AutoModelForCausalLM
 from langchain.llms import HuggingFaceHub
-from transformers import AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
 # from dotenv import load_dotenv
 
 # load_dotenv()
 
-os.environ['HUGGINGFACEHUB_API_TOKEN'] = os.getenv("HF_KEY")
 
 
 def get_vector_store_from_url(url):
@@ -85,11 +84,18 @@ def get_response(user_input):
     #                 lib="avx2", # for CPU
     # )
 
-    llm_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-    llm = HuggingFaceHub(
-        repo_id=llm_model, 
-        model_kwargs={"temperature": 0.3, "max_new_tokens": 250, "top_k": 3}
-    )
+    model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+    # llm = HuggingFaceHub(
+    #     repo_id=llm_model, 
+    #     model_kwargs={"temperature": 0.3, "max_new_tokens": 250, "top_k": 3}
+    # )
+
+    model = transformers.AutoModelForCausalLM.from_pretrained(
+                model_name,
+                trust_remote_code=True,
+                torch_dtype=torch.bfloat16,
+                device_map='auto'
+        
     retriever_chain = get_context_retriever_chain(st.session_state.vector_store,llm)
     conversation_rag_chain = get_conversational_rag_chain(retriever_chain,llm)