commit

README.md

@@ -10,3 +10,54 @@ pinned: false
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# RAG with Gemma
+
+このプロジェクトは、Gemmaモデルを使用したRAG（Retrieval-Augmented Generation）システムの実装です。
+
+## 特徴
+
+- Gemma-2b-itモデルを使用
+- CPUで動作
+- ChromaDBを使用したベクトルストア
+- 日本語対応
+
+## セットアップ
+
+1. 必要なパッケージのインストール:
+```bash
+pip install -r requirements.txt
+```
+
+2. モデルのダウンロード:
+```bash
+python -c "from transformers import AutoTokenizer, AutoModelForCausalLM; AutoTokenizer.from_pretrained('google/gemma-2b-it'); AutoModelForCausalLM.from_pretrained('google/gemma-2b-it')"
+```
+
+## 使用方法
+
+```python
+from rag_system import RAGSystem
+
+# RAGシステムの初期化
+rag = RAGSystem()
+
+# ドキュメントの追加
+documents = [
+    "ドキュメント1の内容",
+    "ドキュメント2の内容",
+    # ...
+]
+rag.add_documents(documents)
+
+# 質問と回答
+question = "質問内容"
+answer = rag.query(question)
+print(f"回答: {answer}")
+```
+
+## 注意事項
+
+- 初回実行時はモデルのダウンロードに時間がかかります
+- CPUでの実行のため、生成に時間がかかる場合があります
+- メモリ使用量に注意してください

rag_system.py

@@ -0,0 +1,103 @@
+import os
+from typing import List
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from sentence_transformers import SentenceTransformer
+import chromadb
+from chromadb.config import Settings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+import torch
+
+class RAGSystem:
+    def __init__(self, model_name: str = "google/gemma-2b-it"):
+        # トークナイザーとモデルの初期化
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float32,
+            device_map="cpu"
+        )
+        
+        # 埋め込みモデルの初期化
+        self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+        
+        # ChromaDBの初期化
+        self.chroma_client = chromadb.Client(Settings(
+            chroma_db_impl="duckdb+parquet",
+            persist_directory="db"
+        ))
+        self.collection = self.chroma_client.get_or_create_collection("documents")
+        
+        # テキスト分割器の初期化
+        self.text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=500,
+            chunk_overlap=50
+        )
+
+    def add_documents(self, documents: List[str]):
+        """ドキュメントをベクトルストアに追加"""
+        chunks = []
+        for doc in documents:
+            chunks.extend(self.text_splitter.split_text(doc))
+        
+        embeddings = self.embedding_model.encode(chunks)
+        
+        # ChromaDBに保存
+        self.collection.add(
+            embeddings=embeddings.tolist(),
+            documents=chunks,
+            ids=[f"doc_{i}" for i in range(len(chunks))]
+        )
+
+    def query(self, question: str, k: int = 3) -> str:
+        """質問に対する回答を生成"""
+        # 質問の埋め込みを取得
+        query_embedding = self.embedding_model.encode(question)
+        
+        # 関連ドキュメントを検索
+        results = self.collection.query(
+            query_embeddings=[query_embedding.tolist()],
+            n_results=k
+        )
+        
+        # コンテキストの構築
+        context = "\n".join(results['documents'][0])
+        
+        # プロンプトの構築
+        prompt = f"""以下のコンテキストに基づいて質問に答えてください。
+
+コンテキスト:
+{context}
+
+質問: {question}
+
+回答:"""
+        
+        # 回答の生成
+        inputs = self.tokenizer(prompt, return_tensors="pt")
+        outputs = self.model.generate(
+            **inputs,
+            max_length=1000,
+            num_return_sequences=1,
+            temperature=0.7
+        )
+        
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+if __name__ == "__main__":
+    # 使用例
+    rag = RAGSystem()
+    
+    # サンプルドキュメントの追加
+    documents = [
+        "RAG（Retrieval-Augmented Generation）は、大規模言語モデルに外部知識を組み込む手法です。",
+        "RAGは、検索と生成を組み合わせることで、より正確な回答を生成することができます。",
+        "RAGの主な利点は、モデルの知識を超えた情報を提供できることです。"
+    ]
+    
+    rag.add_documents(documents)
+    
+    # 質問の例
+    question = "RAGとは何ですか？"
+    answer = rag.query(question)
+    print(f"質問: {question}")
+    print(f"回答: {answer}") 

requirements.txt

@@ -0,0 +1,7 @@
+transformers>=4.38.0
+sentence-transformers>=2.2.2
+faiss-cpu>=1.7.4
+langchain>=0.1.0
+chromadb>=0.4.22
+tqdm>=4.66.1
+python-dotenv>=1.0.0 

test.py

@@ -1 +0,0 @@
-print("Hello, World!")