Update app.py

app.py

@@ -1,21 +1,20 @@
 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.sentence_transformer import SentenceTransformerEmbeddings
+from langchain.embeddings import HuggingFaceHubEmbeddings
+from langchain.llms import HuggingFaceHub
 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():
@@ -47,28 +46,14 @@ 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 = SentenceTransformerEmbeddings(model_name="jhgan/ko-sroberta-multitask")
+    embeddings = HuggingFaceHubEmbeddings(model_name="jhgan/ko-sroberta-multitask", huggingfacehub_api_token=HF_TOKEN)
 
     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="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)
+    #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)
     
     global qa
     qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, return_source_documents=True)