import gradio as gr
"""
import pathlib
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
"""
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import pandas as pd

## LLMの読み込み（Qwen2.5-3Bをsafetensorsで読み込み）
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "Qwen/Qwen2.5-3B"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
llm = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16)
llm.eval()

"""
## LLMの読み込み
models_dir = pathlib.Path(__file__).parent / "models"
models_dir.mkdir(exist_ok=True)

model_path = hf_hub_download(
    repo_id="Mori-kamiyama/sarashina2-13b-r1",
    filename="model.gguf",
    local_dir=models_dir
)

llm = Llama(model_path=model_path)
"""

## 埋め込みモデルの読み込み
model = SentenceTransformer("BAAI/bge-m3")

# ドキュメントの読み込み
df = pd.read_csv("document.csv")
# "text"カラムをリストとして抽出
texts = df['text'].tolist()

# ベクトル化
doc_embeddings = model.encode(texts, normalize_embeddings=True)

# FAISSのセットアップ
dimension = doc_embeddings.shape[1]
index = faiss.IndexFlatIP(dimension)  # Cosine用にnormalize済ならこれ
index.add(np.array(doc_embeddings))

def generate_text(prompt):
    result = llm(search(prompt))
    return result['choices'][0]['text']

def search(query):
    query_embedding = model.encode([query], normalize_embeddings=True)

    # FAISSで検索
    top_k = 2
    D, I = index.search(np.array(query_embedding), top_k)

    retrieved_docs = []
    print("\n🔍 検索結果:")
    for idx in I[0]:
        doc_text = texts[idx]
        retrieved_docs.append(doc_text)
        print(f"→ {doc_text}")

    # RAG用のプロンプトを作成
    prompt = "以下の文書を参照して質問に答えてください。\n\n文書:\n"
    prompt += "\n".join(retrieved_docs)
    prompt += f"\n\n質問: {query}"

    return prompt

iface = gr.Interface(fn=generate_text,
                     inputs="text",
                     outputs="text",
                     title="sarashina-R13B-RAG")

iface.launch()