Create app.py

app.py

@@ -0,0 +1,58 @@
+import streamlit as st
+import pandas as pd
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+
+# Define paths for the dataset splits
+splits = {
+    'train': 'data/train-00000-of-00001.parquet',
+    'validation': 'data/validation-00000-of-00001.parquet',
+    'test': 'data/test-00000-of-00001.parquet'
+}
+
+# Load the dataset
+@st.cache_resource
+def load_dataset(split="train"):
+    return pd.read_parquet(f"hf://datasets/BEE-spoke-data/survivorslib-law-books/{splits[split]}")
+
+# Initialize the model and tokenizer
+@st.cache_resource
+def load_model():
+    model_name = "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+# Streamlit interface
+st.title("Legal Text Generator with NVIDIA Llama")
+st.write("Generate text based on the Survivorslib Legal Dataset and the NVIDIA Llama model.")
+
+# Load dataset and model pipeline
+st.sidebar.title("Options")
+split_option = st.sidebar.selectbox("Select dataset split", ["train", "validation", "test"])
+dataset = load_dataset(split=split_option)
+text_generator = load_model()
+
+# Show sample data from the dataset
+st.subheader(f"Sample Data from {split_option.capitalize()} Split")
+st.write(dataset.head())  # Displaying the first few rows of the selected dataset split
+
+# Prompt input
+prompt = st.text_area("Enter your prompt:", placeholder="Type a legal prompt or select a sample text...")
+
+# Optional: Select sample text from the dataset to use as a prompt
+if st.button("Use Sample Text"):
+    if 'content' in dataset.columns:
+        prompt = dataset['content'].iloc[0]
+        st.write(f"Using sample text from dataset: {prompt}")
+    else:
+        st.write("Dataset does not contain a 'content' column with text data.")
+
+# Generate text based on the prompt
+if st.button("Generate Response"):
+    if prompt:
+        with st.spinner("Generating response..."):
+            generated_text = text_generator(prompt, max_length=100, do_sample=True, temperature=0.7)[0]["generated_text"]
+            st.write("**Generated Text:**")
+            st.write(generated_text)
+    else:
+        st.write("Please enter a prompt to generate a response.")