Update app.py

app.py

@@ -1,16 +1,16 @@
-import os
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 from datasets import load_dataset
 import faiss
 import gradio as gr
 from accelerate import Accelerator
+import os
+import torch
 
 # 환경 변수에서 Hugging Face API 키 로드
 hf_api_key = os.getenv('HF_API_KEY')
 
-# 모델 ID 및 토크나이저 설정
+# 모델 및 토크나이저 설정
 model_id = "microsoft/phi-2"
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_api_key, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
@@ -20,11 +20,9 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.float32
 )
 
-# Accelerator 설정
 accelerator = Accelerator()
 model = accelerator.prepare(model)
 
-# 데이터셋 및 FAISS 인덱스 로드
 ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
 dataset = load_dataset("not-lain/wikipedia", revision="embedded")
 data = dataset["train"]
@@ -44,7 +42,14 @@ def format_prompt(prompt, retrieved_documents, k):
 def generate(formatted_prompt):
     prompt_text = f"{SYS_PROMPT} {formatted_prompt}"
     input_ids = tokenizer(prompt_text, return_tensors="pt", padding=True).input_ids.to(accelerator.device)
-    outputs = model.generate(input_ids, max_new_tokens=1024, eos_token_id=tokenizer.eos_token_id, do_sample=True, temperature=0.6, top_p=0.9)
+    outputs = model.generate(
+        input_ids,
+        max_new_tokens=1024,
+        eos_token_id=tokenizer.eos_token_id,
+        do_sample=True,
+        temperature=0.6,
+        top_p=0.9
+    )
     return tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
 
 def rag_chatbot_interface(prompt: str, k: int = 2):
@@ -59,8 +64,7 @@ iface = gr.Interface(
     inputs="text",
     outputs="text",
     title="Retrieval-Augmented Generation Chatbot",
-    description="This chatbot provides more accurate answers by searching relevant documents and generating responses.",
-    share=True  # 공개 링크 생성
+    description="This chatbot provides more accurate answers by searching relevant documents and generating responses."
 )
 
-iface.launch()
+iface.launch(share=True)  # 여기에서 share=True를 설정하여 공개 링크를 생성