Update app.py

app.py

@@ -1,20 +1,21 @@
 import gradio as gr
 import pandas as pd
 import json
-import os
 
 from langchain.document_loaders import DataFrameLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.embeddings import HuggingFaceHubEmbeddings
-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
 from trafilatura.settings import use_config
 
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
 
 
 def loading_website():
@@ -46,14 +47,28 @@ def url_changes(url, pages_to_visit, urls_to_scrape, repo_id):
     texts = text_splitter.split_documents(documents)
     print(f"documents splitted into {len(texts)} chunks") 
     
-    embeddings = HuggingFaceHubEmbeddings(model_name="jhgan/ko-sroberta-multitask", huggingfacehub_api_token=HF_TOKEN)
+    embeddings = SentenceTransformerEmbeddings(model_name="jhgan/ko-sroberta-multitask")
 
     persist_directory = './vector_db'
     db = Chroma.from_documents(texts, embeddings, persist_directory=persist_directory)
     retriever = db.as_retriever()
 
-    #MODEL = 'beomi/KoAlpaca-Polyglot-5.8B'
-    llm = HuggingFaceHub(repo_id="beomi/KoAlpaca-Polyglot-5.8B", model_kwargs={"temperature":0.6, "max_new_tokens":250}, huggingfacehub_api_token=HF_TOKEN)
+    MODEL = 'beomi/KoAlpaca-Polyglot-5.8B'
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL,
+        torch_dtype="auto",
+    )
+    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)
     
     global qa
     qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, return_source_documents=True)