Update app.py

app.py

@@ -4,7 +4,6 @@ import pickle
 from langchain.prompts import ChatPromptTemplate
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_huggingface import HuggingFaceEmbeddings
-from transformers import pipeline
 from langchain_community.llms import HuggingFacePipeline
 from langchain.retrievers import ParentDocumentRetriever
 from langchain.storage import InMemoryStore
@@ -14,18 +13,18 @@ from langchain_core.prompts import ChatPromptTemplate, FewShotChatMessagePromptT
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnableLambda
 from datetime import date
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 # from setup import download_olmo_model, OLMO_MODEL
 
 # Ensure model is downloaded before proceeding
-@st.cache_resource
-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
+# @st.cache_resource
+# 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, tokenizer = load_model()
+# model, tokenizer = load_model()
 
 # # Define the path to your bash script
 # script_path = "./start.sh"
@@ -96,30 +95,20 @@ 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, 
-                    tokenizer=tokenizer, 
-                    max_length=4000,
-                    temperature=temperature,
-                    top_p=0.95,
-                    repetition_penalty=1.15
-                )
+            "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
+        )
 
-    llm = HuggingFacePipeline(pipeline=pipe)
 
-    
+    llm = HuggingFacePipeline(pipeline=pipe)
 
     
-    # Initialize LangChain
-    # llm = HuggingFaceLLM(
-    #         model_id="EleutherAI/gpt-neo-1.3B",  # or another suitable model
-    #         temperature=temperature,
-    #         max_tokens=256
-    #     )    
-    
     today = date.today()
     # Response prompt 
     response_prompt_template = """You are an assistant who helps Ocean Hack Week community to answer their questions. I am going to ask you a question. Your response should be comprehensive and not contradicted with the following context if they are relevant. Otherwise, ignore them if they are not relevant.