Update app.py

app.py

@@ -1,23 +1,15 @@
+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')
-
-# 모델 및 토크나이저 설정
 model_id = "microsoft/phi-2"
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_api_key, trust_remote_code=True)
-
-# 토크나이저에 패딩 토큰 설정
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token  # EOS 토큰을 패딩 토큰으로 사용
-
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     token=hf_api_key,
@@ -44,19 +36,6 @@ def format_prompt(prompt, retrieved_documents, k):
         PROMPT += f"{retrieved_documents['text'][idx]}\n"
     return PROMPT
 
-def generate(formatted_prompt):
-    prompt_text = f"{SYS_PROMPT} {formatted_prompt}"
-    input_ids = tokenizer(prompt_text, return_tensors="pt", padding="max_length", max_length=512).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
-    )
-    return tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
-
 def rag_chatbot_interface(prompt: str, k: int = 2):
     scores, retrieved_documents = search(prompt, k)
     formatted_prompt = format_prompt(prompt, retrieved_documents, k)