Zee292/retrieval_augmented_generation

1import streamlit as st2import numpy as np3import faiss4from sentence_transformers import SentenceTransformer5from groq import Groq6import os7 8# config9st.set_page_config(10    page_title="rag itera - kelompok 4",11    page_icon="🧠",12    layout="centered"13)14 15st.title("nlp kelompok 4")16st.subheader("tugas b - rag sistem tanya jawab itera")17 18# api key19try:20    groq_api_key = st.secrets["GROQ_API_KEY"]21except:22    groq_api_key = os.getenv("GROQ_API_KEY")23 24if not groq_api_key:25    st.error("❌ API key Groq tidak ditemukan.")26    st.stop()27 28client = Groq(29    api_key=groq_api_key30)31# dokumen32documents = [33    "Institut Teknologi Sumatera atau ITERA berada di Lampung Selatan.",34    "ITERA didirikan untuk meningkatkan kualitas pendidikan teknologi di Pulau Sumatera.",35    "Fakultas di ITERA terdiri dari FTIK, FTI, dan FTSPK.",36    "Program studi Teknik Informatika berada di bawah Fakultas Teknologi Industri.",37    "Perpustakaan ITERA menyediakan layanan peminjaman buku dan akses jurnal digital.",38    "Masjid At Tanwir merupakan pusat kegiatan keagamaan di ITERA.",39    "Gedung Kuliah Umum ITERA sering digunakan untuk seminar nasional.",40    "Mahasiswa ITERA aktif mengikuti kompetisi nasional dan internasional.",41    "ITERA memiliki berbagai organisasi mahasiswa dan unit kegiatan mahasiswa.",42    "Kampus ITERA memiliki luas lebih dari 200 hektar."43]44 45# load model & faiss46@st.cache_resource47def load_retrieval():48    model = SentenceTransformer('all-MiniLM-L6-v2')49    doc_embeddings = model.encode(documents)50    dimension = doc_embeddings.shape[1]51    index = faiss.IndexFlatL2(dimension)52    index.add(np.array(doc_embeddings))53    return model, index54 55# fungsi retrieval56def retrieve(query, k=3):57    query_embedding = st.session_state['model'].encode([query])58    distances, indices = st.session_state['index'].search(59        np.array(query_embedding),60        k61    )62    results = []63    for idx in indices[0]:64        results.append(documents[idx])65    return results, indices[0]66 67# fungsi rag68def ask_rag(query, k=3):69    retrieved_docs, indices = retrieve(query, k)70    context = "\n".join(retrieved_docs)71    prompt = f"""72    jawab pertanyaan berdasarkan konteks berikut.73 74    konteks:75    {context}76 77    pertanyaan:78    {query}79 80    jika jawaban tidak ada di konteks, katakan "tidak ditemukan".81    """82    response = client.chat.completions.create(83        model="llama-3.3-70b-versatile",84        messages=[{"role": "user", "content": prompt}],85        temperature=0.386    )87    answer = response.choices[0].message.content88    token_usage = response.usage.total_tokens89    return answer, retrieved_docs, token_usage90 91# inisialisasi92if st.button("inisialisasi sistem rag", use_container_width=True):93    with st.spinner("memuat sentence transformer dan faiss..."):94        model, index = load_retrieval()95        st.session_state['model'] = model96        st.session_state['index'] = index97    st.success("sistem rag berhasil diinisialisasi!")98 99# interface100if 'model' in st.session_state:101    st.markdown("---")102    query = st.text_input(103        "masukkan pertanyaan:",104        "apa saja fakultas yang ada di itera?"105    )106    k = st.slider(107        "jumlah dokumen retrieval",108        min_value=1,109        max_value=5,110        value=3111    )112    if st.button("cari & jawab", use_container_width=True):113        with st.spinner("mencari dokumen relevan..."):114            answer, docs, tokens = ask_rag(query, k)115        st.subheader("dokumen relevan")116        for i, doc in enumerate(docs, 1):117            st.info(f"{i}. {doc}")118        st.subheader("jawaban ai")119        st.success(answer)120        st.caption(f"total token usage: {tokens}")121 122@st.cache_resource123def load_retrieval():124 125    model = SentenceTransformer('all-MiniLM-L6-v2')126 127    doc_embeddings = model.encode(documents)128 129    dimension = doc_embeddings.shape[1]130 131    index = faiss.IndexFlatL2(dimension)132 133    index.add(np.array(doc_embeddings))134 135    return model, index136