Update app.py

app.py

@@ -14,6 +14,8 @@ from langchain_community.document_loaders import PyPDFLoader
 from langchain_chroma import Chroma
 from langchain_community.vectorstores import Chroma
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import transformers
+import torch
 
 #st.set_page_config(page_title="Document Genie", layout="wide")
 
@@ -81,8 +83,17 @@ def get_conversational_chain():
     #repo_id='meta-llama/Meta-Llama-3-70B'
     #llm = HuggingFaceEndpoint(
     #repo_id=repo_id, max_length=512, temperature=0.5, token=HUGGING_FACE_API_KEY)
-    tokenizer = AutoTokenizer.from_pretrained("google/gemma-1.1-2b-it")
-    llm = AutoModelForCausalLM.from_pretrained("google/gemma-1.1-2b-it")
+    #tokenizer = AutoTokenizer.from_pretrained("google/gemma-1.1-2b-it")
+    #llm = AutoModelForCausalLM.from_pretrained("google/gemma-1.1-2b-it")
+
+    model_id = "google/gemma-1.1-2b-it"
+    dtype = torch.bfloat16
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    llm= AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=dtype,
+    )
 
     pt = ChatPromptTemplate.from_template(template)
     # Retrieve and generate using the relevant snippets of the blog.