import os import io import torch import gradio as gr from typing import List, Tuple, Dict from sentence_transformers import SentenceTransformer from transformers import AutoProcessor, Gemma3ForConditionalGeneration from pypdf import PdfReader # -------------------------------------------------------------------- # Konfiguration # -------------------------------------------------------------------- EMBED_MODEL_ID = "google/embeddinggemma-300m" LLM_MODEL_ID = "google/gemma-3-4b-it" # Globale States (simpler, reicht für einen Space) EMBED_MODEL: SentenceTransformer = None LLM_MODEL: Gemma3ForConditionalGeneration = None LLM_PROCESSOR: AutoProcessor = None DOCUMENT_STORE: List[Dict] = [] # {content, embedding (tensor), source} # -------------------------------------------------------------------- # Model Loading # -------------------------------------------------------------------- def get_device() -> torch.device: if torch.cuda.is_available(): return torch.device("cuda") return torch.device("cpu") def get_embedding_model() -> SentenceTransformer: global EMBED_MODEL if EMBED_MODEL is None: device = "cuda" if torch.cuda.is_available() else "cpu" # EmbeddingGemma nutzt SentenceTransformers-Wrapper EMBED_MODEL = SentenceTransformer(EMBED_MODEL_ID, device=device) return EMBED_MODEL def get_llm() -> Tuple[Gemma3ForConditionalGeneration, AutoProcessor]: global LLM_MODEL, LLM_PROCESSOR if LLM_MODEL is None or LLM_PROCESSOR is None: device = get_device() if device.type == "cuda": # bfloat16 für GPU wie im Model-Card-Beispiel LLM_MODEL = Gemma3ForConditionalGeneration.from_pretrained( LLM_MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto", ).eval() else: # CPU: float32 LLM_MODEL = Gemma3ForConditionalGeneration.from_pretrained( LLM_MODEL_ID, torch_dtype=torch.float32, ).to(device).eval() LLM_PROCESSOR = AutoProcessor.from_pretrained(LLM_MODEL_ID) return LLM_MODEL, LLM_PROCESSOR # -------------------------------------------------------------------- # Datei-Handling & Chunking # -------------------------------------------------------------------- def extract_text_from_file(path: str) -> str: ext = os.path.splitext(path)[1].lower() if ext in [".txt", ".md", ".markdown"]: with open(path, "r", encoding="utf-8", errors="ignore") as f: return f.read() if ext == ".pdf": text = [] reader = PdfReader(path) for page in reader.pages: page_text = page.extract_text() if page_text: text.append(page_text) return "\n".join(text) # Fallback: Versuche als Text zu lesen try: with open(path, "r", encoding="utf-8", errors="ignore") as f: return f.read() except Exception: return "" def chunk_text( text: str, chunk_size: int = 800, # ca. "Token"-Näherung über Wörter chunk_overlap: int = 200, ) -> List[str]: words = text.split() if not words: return [] chunks = [] start = 0 while start < len(words): end = min(len(words), start + chunk_size) chunk = " ".join(words[start:end]).strip() if chunk: chunks.append(chunk) if end == len(words): break start = max(0, end - chunk_overlap) return chunks # -------------------------------------------------------------------- # Indexing / RAG # -------------------------------------------------------------------- def index_files(file_paths: List[str]) -> str: """ Liest hochgeladene Dateien ein, chunked sie und legt Embeddings in einem einfachen In-Memory-Store ab. """ if not file_paths: return "Keine Dateien übergeben." embed_model = get_embedding_model() added_chunks = 0 for path in file_paths: if path is None: continue text = extract_text_from_file(path) if not text.strip(): continue chunks = chunk_text(text) if not chunks: continue # Embeddings mit EmbeddingGemma – nutzt automatisch passende Prompts doc_embeddings = embed_model.encode_document( chunks, convert_to_tensor=True, ) # shape: (num_chunks, dim) for chunk, emb in zip(chunks, doc_embeddings): DOCUMENT_STORE.append( { "content": chunk, "embedding": emb, # torch.Tensor "source": os.path.basename(path), } ) added_chunks += 1 if added_chunks == 0: return "Keine verwertbaren Text-Chunks gefunden." return f"Index aktualisiert: {len(DOCUMENT_STORE)} Chunks gespeichert (neu hinzugefügt: {added_chunks})." def clear_index() -> str: DOCUMENT_STORE.clear() return "Index geleert (0 Chunks)." def retrieve_relevant_chunks( query: str, top_k: int = 5, ) -> List[Dict]: """ Einfacher Dense-Retriever auf Basis von EmbeddingGemma-300M. Nutzt die SentenceTransformers-Similarity-Funktion (inner product / cosine). """ if not DOCUMENT_STORE: return [] embed_model = get_embedding_model() # Query-Embedding q_emb = embed_model.encode_query( query, convert_to_tensor=True, ) # (dim,) # Stack aller Dokument-Embeddings doc_embs = torch.stack([d["embedding"] for d in DOCUMENT_STORE]) # (N, dim) # Similarity via SentenceTransformers-API (richtige Metrik laut Config) sims = embed_model.similarity(q_emb, doc_embs)[0] # (N,) top_k = min(top_k, len(DOCUMENT_STORE)) scores, indices = torch.topk(sims, k=top_k) results = [] for score, idx in zip(scores.tolist(), indices.tolist()): entry = DOCUMENT_STORE[idx] results.append( { "content": entry["content"], "source": entry["source"], "score": float(score), } ) return results # -------------------------------------------------------------------- # LLM-Generierung (Gemma-3-4B-IT) # -------------------------------------------------------------------- def build_system_prompt() -> str: return ( "Du bist ein hilfreicher, präziser Assistent. " "Du beantwortest Fragen basierend auf bereitgestellten Dokumenten. " "Wenn Informationen nicht im Kontext sind, sag explizit, dass sie nicht im Kontext stehen. " "Antworte bitte auf Deutsch und zitiere ggf. die relevanten Teile in eigenen Worten." ) def build_user_prompt(user_question: str, retrieved_chunks: List[Dict]) -> str: if not retrieved_chunks: return ( f"Benutzerfrage:\n{user_question}\n\n" "Es wurden keine Dokumente im Kontext gefunden. " "Beantworte die Frage so gut wie möglich, aber weise darauf hin, " "dass du keinen Dokumentenkontext hast." ) context_str_parts = [] for i, ch in enumerate(retrieved_chunks, start=1): context_str_parts.append( f"[{i}] (Quelle: {ch['source']}, Score: {ch['score']:.3f})\n{ch['content']}" ) context_str = "\n\n".join(context_str_parts) prompt = ( f"Benutzerfrage:\n{user_question}\n\n" "Hier sind die relevantesten Kontextpassagen aus den hochgeladenen Dokumenten:\n" f"{context_str}\n\n" "Aufgabe:\n" "- Nutze primär diese Kontexte.\n" "- Wenn etwas nicht im Kontext steht, sag das klar.\n" "- Fasse die relevanten Stellen strukturiert zusammen.\n" "- Antworte auf Deutsch.\n\n" "Antwort:" ) return prompt def answer_with_rag(question: str, history: List[Tuple[str, str]]) -> str: model, processor = get_llm() # 1. Retrieve relevante Chunks retrieved = retrieve_relevant_chunks(question, top_k=5) # 2. Prompt bauen system_prompt = build_system_prompt() user_prompt = build_user_prompt(question, retrieved) messages = [ { "role": "system", "content": [{"type": "text", "text": system_prompt}], }, { "role": "user", "content": [{"type": "text", "text": user_prompt}], }, ] # 3. Chat Template + Generation inputs = processor.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt", ).to(model.device, dtype=model.dtype) input_len = inputs["input_ids"].shape[-1] with torch.inference_mode(): generated = model.generate( **inputs, max_new_tokens=512, do_sample=True, temperature=0.7, top_p=0.9, ) gen_tokens = generated[0][input_len:] answer = processor.decode(gen_tokens, skip_special_tokens=True) return answer # -------------------------------------------------------------------- # Gradio UI # -------------------------------------------------------------------- def ingest_files_ui(files) -> str: if not files: return "Bitte zuerst eine oder mehrere Dateien hochladen." paths = [f.name if hasattr(f, "name") else f for f in files] return index_files(paths) def chat_fn(message: str, history: List[Tuple[str, str]]): """ Callback für gr.Chatbot: nimmt die neue User-Nachricht, macht RAG+LLM und hängt Antwort an die History an. """ if not message or not message.strip(): return history answer = answer_with_rag(message, history) history = history + [(message, answer)] return history def build_demo() -> gr.Blocks: with gr.Blocks(title="Gemma 3 RAG – EmbeddingGemma-300M", theme="soft") as demo: gr.Markdown( """ # 🔍 Gemma 3 RAG Space **RAG-Pipeline mit `google/embeddinggemma-300m` + `google/gemma-3-4b-it`** Schritte: 1. Links Dateien hochladen und indexieren. 2. Rechts im Chat Fragen zu deinen Dokumenten stellen. """ ) with gr.Row(): with gr.Column(scale=1): gr.Markdown("### 📁 Dokumente") file_uploader = gr.File( label="Dateien hochladen (.pdf, .txt, .md, ...)", file_count="multiple", type="filepath", ) index_button = gr.Button("🔄 Index aktualisieren") clear_index_button = gr.Button("🧹 Index leeren") index_status = gr.Markdown("Noch keine Dokumente indexiert.") index_button.click( fn=index_files, inputs=file_uploader, outputs=index_status, ) clear_index_button.click( fn=clear_index, inputs=None, outputs=index_status, ) with gr.Column(scale=2): gr.Markdown("### 💬 Chat mit Gemma-3 über deine Dateien") chatbot = gr.Chatbot( label="Kontext-sensitiver Chat", type="tuple", show_copy_button=True, ) msg = gr.Textbox( label="Deine Frage", placeholder="Frag etwas zu den hochgeladenen Dokumenten...", lines=2, ) send_btn = gr.Button("Senden") clear_chat_btn = gr.Button("Chat löschen") def user_submit(user_message, chat_history): if not user_message: return "", chat_history chat_history = chat_history + [(user_message, None)] return "", chat_history # Ablauf: # 1. User-Text -> History mit (msg, None) # 2. Chat-Funktion ersetzt None durch generierte Antwort msg.submit(user_submit, [msg, chatbot], [msg, chatbot]).then( chat_fn, chatbot, chatbot ) send_btn.click(user_submit, [msg, chatbot], [msg, chatbot]).then( chat_fn, chatbot, chatbot ) clear_chat_btn.click( fn=lambda: [], inputs=None, outputs=chatbot, ) return demo demo = build_demo() if __name__ == "__main__": # In HF Spaces wird normalerweise automatisch gestartet, # aber lokal brauchst du das: demo.launch()