Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
 from langchain import PromptTemplate, HuggingFaceHub, LLMChain
 import os
 from langchain.memory import ConversationBufferMemory
@@ -19,10 +20,10 @@ prompt = PromptTemplate(template=template, input_variables=["question"])
 
 answer_memory = ConversationBufferMemory(input_key='question', memory_key='chat_history')
 
+tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-base-alpha-3b")
 
-llm = HuggingFaceHub(repo_id="stabilityai/stablelm-tuned-alpha-3b" , model_kwargs={"temperature":0, "max_length":64})
-
-llm_chain = LLMChain(prompt=prompt, llm=llm, verbose=True, output_key='answer', memory=answer_memory)
+model = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-base-alpha-3b")
+llm_chain = LLMChain(prompt=prompt, llm=model, verbose=True, output_key='answer', memory=answer_memory)
 
 
 wiki = WikipediaAPIWrapper()