app.py

import streamlit as st
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
import openai

# Load FAISS index
index = faiss.read_index("faiss_index.bin")

# Load embedding model
model = SentenceTransformer("all-MiniLM-L6-v2")

# OpenAI API Key (store it as a secret in Hugging Face)
openai.api_key = st.secrets["GROQ_API_KEY"]

# Load the preprocessed Pile Law dataset (replace with your dataset path)
law_data = ["Sample Legal Document 1...", "Sample Legal Document 2..."]  # Replace with actual data loading

# Function to search relevant legal documents
def search_legal_docs(query, top_k=5):
    query_embedding = model.encode([query])
    _, idxs = index.search(query_embedding, top_k)
    return [law_data[i] for i in idxs[0]]

# Streamlit UI
st.title("🔍 Legal AI Assistant (Pile Law)")

query = st.text_input("📌 Enter your legal query:")

if query:
    results = search_legal_docs(query)
    st.write("### 📄 Relevant Legal Documents:")
    for res in results:
        st.write(f"- {res}")

    # Generate AI-based legal response
    response = openai.ChatCompletion.create(
        model="gpt-4",
        messages=[{"role": "system", "content": "You are a legal assistant."},
                  {"role": "user", "content": query}]
    )
    st.write("### 🧑‍⚖️ AI Response:")
    st.write(response['choices'][0]['message']['content'])