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-r2",
    filename="model.gguf",
    local_dir=models_dir
)

llm = Llama(model_path=model_path)

## 埋め込みモデルの読み込み
model = SentenceTransformer("BAAI/bge-small-en-v1.5")

# ドキュメントの読み込み
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):
    full_prompt = search(prompt)
    output = llm(full_prompt, max_tokens=128)
    return output["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)

    # スコアが低すぎる場合は結果を無視
    if D[0][0] < 0.3:
        print("🔍 十分な類似度のある検索結果が見つかりませんでした。")
        return "関連する情報はありませんでした。そのまま回答してください。"

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

    # RAG用のプロンプトを作成
    prompt = "以下の文書を参照して質問に答えてください。"
    prompt += "生成は以下のフォーマットで回答しなさい\n<reasoning>\n...\n</reasoning>\n\n<answer>\n...\n</answer>"
    prompt += "\n\n文書:\n"
    prompt += "\n".join(retrieved_docs)
    prompt += "\n必ずこの質問だけ一回だけ回答してください。文体はやさしく柔らかくお願いします。"
    prompt += f"\n\n質問: {query}"

    return prompt

iface = gr.Interface(fn=generate_text,
                     inputs="text",
                     outputs="text",
                     title="sarashina2-13b-r2",)

iface.launch()