Update app.py

app.py

@@ -25,6 +25,25 @@ dataset = load_dataset("not-lain/wikipedia", revision="embedded")
 data = dataset["train"]
 data = data.add_faiss_index("embeddings")
 
+def generate(formatted_prompt):
+    prompt_text = f"{SYS_PROMPT} {formatted_prompt}"
+    # 토크나이징 시 attention_mask도 함께 생성
+    encoding = tokenizer(prompt_text, return_tensors="pt", padding="max_length", max_length=512, truncation=True)
+    input_ids = encoding['input_ids'].to(accelerator.device)
+    attention_mask = encoding['attention_mask'].to(accelerator.device)
+
+    outputs = model.generate(
+        input_ids,
+        attention_mask=attention_mask,  # attention_mask 전달
+        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], skip_special_tokens=True)
+
+
 def search(query: str, k: int = 3):
     embedded_query = ST.encode(query)
     scores, retrieved_examples = data.get_nearest_examples("embeddings", embedded_query, k=k)