Update app.py

app.py

@@ -4,10 +4,12 @@ import json
 
 from langchain.document_loaders import DataFrameLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.llms import HuggingFaceHub
 from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
 from langchain.vectorstores import Chroma
 from langchain.chains import RetrievalQA
+from langchain import HuggingFacePipeline
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 from trafilatura import fetch_url, extract
 from trafilatura.spider import focused_crawler
@@ -50,8 +52,25 @@ def url_changes(url, pages_to_visit, urls_to_scrape, repo_id):
     persist_directory = './vector_db'
     db = Chroma.from_documents(texts, embeddings, persist_directory=persist_directory)
     retriever = db.as_retriever()
+
+    MODEL = 'beomi/KoAlpaca-Polyglot-5.8B'
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL,
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True,
+    ).to(device=f"cuda", non_blocking=True)
+    model.eval()
+    pipe = pipeline(
+        'text-generation',
+        model=model,
+        tokenizer=MODEL,
+        max_length=512,
+        temperature=0,
+        top_p=0.95,
+        repetition_penalty=1.15
+    )
+    llm = HuggingFacePipeline(pipeline=pipe)
     
-    llm = HuggingFaceHub(repo_id=repo_id, model_kwargs={"temperature":0.1, "max_new_tokens":250})
     global qa
     qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, return_source_documents=True)
     return "Ready"