Update app.py

app.py

@@ -21,7 +21,12 @@ os.environ['LANGCHAIN_ENDPOINT'] = 'https://api.smith.langchain.com'
 os.environ['LANGCHAIN_API_KEY'] = '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'
 
 
-
+@st.cache_resource
+def load_model():
+    model_name = "bigscience/bloom-1b7"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+    return model, tokenizer
 def load_from_pickle(filename):
     with open(filename, "rb") as file:
         return pickle.load(file)
@@ -70,16 +75,17 @@ def get_chain(temperature):
     child_splitter = RecursiveCharacterTextSplitter(chunk_size=300,
                                                     chunk_overlap=50)
     retriever = load_retriever(docstore_path,chroma_path,embeddings,child_splitter,parent_splitter)
-    model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-    # Replace the local OLMOLLM with the Hugging Face model
-    pipe = pipeline(
-            "text-generation",
-            model=model_name,
-            max_length=3000,  # Allows for 2,093 input tokens + some generated tokens
-            max_new_tokens=500,  # Generates up to 100 new tokens
-            temperature=temperature  # Adjust temperature for response creativity
-        )
+    model, tokenizer = load_model()
 
+    pipe = pipeline(
+        "text-generation",
+        model=model, 
+        tokenizer=tokenizer, 
+        max_length=2048,
+        temperature=temperature,
+        top_p=0.95,
+        repetition_penalty=1.15
+    )
 
     llm = HuggingFacePipeline(pipeline=pipe)