Update app.py

app.py

@@ -15,20 +15,17 @@ from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnableLambda
 from datetime import date
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from setup import download_olmo_model, OLMO_MODEL
+# from setup import download_olmo_model, OLMO_MODEL
 
 # Ensure model is downloaded before proceeding
 @st.cache_resource
-def ensure_model_downloaded():
-    try:
-        model_path = download_olmo_model()
-        st.success(f"Model successfully loaded from {model_path}")
-        return model_path
-    except Exception as e:
-        st.error(f"Failed to download or locate the model: {str(e)}")
-        st.stop()
+def load_model():
+    model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+    return model, tokenizer
 
-model_path = ensure_model_downloaded()# import subprocess
+model, tokenizer = load_model()
 
 # # Define the path to your bash script
 # script_path = "./start.sh"
@@ -101,18 +98,17 @@ def get_chain(temperature):
     retriever = load_retriever(docstore_path,chroma_path,embeddings,child_splitter,parent_splitter)
     
     # Replace the local OLMOLLM with the Hugging Face model
-    try:
-        llm = LlamaCpp(
-            model_path=str(model_path),
-            temperature=temperature,
-            max_tokens=3000,
-            n_ctx=4096,
-            verbose=False,
-            echo=False
-        )
-    except Exception as e:
-        st.error(f"Failed to initialize LlamaCpp: {str(e)}")
-        st.stop()
+    pipe = pipeline(
+                    "text-generation",
+                    model=model, 
+                    tokenizer=tokenizer, 
+                    max_length=4000,
+                    temperature=temperature,
+                    top_p=0.95,
+                    repetition_penalty=1.15
+                )
+
+    llm = HuggingFacePipeline(pipeline=pipe)