|
import gradio as gr |
|
import spaces |
|
import torch |
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
from services.rag_pipeline import rag_pipeline |
|
|
|
model_name = "dasomaru/gemma-3-4bit-it-demo" |
|
|
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) |
|
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
model_name, |
|
torch_dtype=torch.float16, |
|
trust_remote_code=True, |
|
) |
|
|
|
|
|
search_cache = {} |
|
|
|
@spaces.GPU(duration=300) |
|
def generate_response(query: str): |
|
tokenizer = AutoTokenizer.from_pretrained("dasomaru/gemma-3-4bit-it-demo") |
|
model = AutoModelForCausalLM.from_pretrained("dasomaru/gemma-3-4bit-it-demo") |
|
model.to("cuda") |
|
|
|
if query in search_cache: |
|
print(f"⚡ 캐시 사용: '{query}'") |
|
return search_cache[query] |
|
|
|
|
|
results = rag_pipeline(query) |
|
|
|
|
|
if isinstance(results, list): |
|
results = "\n\n".join(results) |
|
|
|
search_cache[query] = results |
|
return results |
|
|
|
|
|
demo = gr.Interface( |
|
fn=generate_response, |
|
inputs=gr.Textbox(lines=2, placeholder="질문을 입력하세요"), |
|
outputs="text", |
|
title="Law RAG Assistant", |
|
description="법령 기반 RAG 파이프라인 테스트", |
|
) |
|
|
|
|
|
demo.launch() |
|
|