Update app.py

app.py

@@ -1,58 +1,87 @@
+# app.py
+
+import gradio as gr
 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)
+from transformers import pipeline
+from datasets import load_dataset
+
+# Initialize text-generation pipeline with the model
+model_name = "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF"
+pipe = pipeline("text-generation", model=model_name)
+
+# Load the dataset from the cloned local directory
+ds = load_dataset("./canadian-legal-data", split="train")
+
+# Gradio Interface setup
+def respond(
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+):
+    messages = [{"role": "system", "content": system_message}]
+    
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
+    
+    messages.append({"role": "user", "content": message})
+
+    response = ""
+    
+    for message in pipe(
+        prompt=message,
+        max_length=max_tokens,
+        do_sample=True,
+        temperature=temperature,
+        top_p=top_p,
+    ):
+        token = message["generated_text"]
+        response += token
+        yield response
+
+# Streamlit Interface setup
+def streamlit_interface():
+    st.title("Canadian Legal Text Generator")
+    st.write("Enter a prompt related to Canadian legal data and generate text using Llama-3.1.")
+    
+    # Show dataset sample
+    st.subheader("Sample Data from Canadian Legal Dataset:")
+    st.write(ds[:5])  # Display the first 5 rows of the dataset
+
+    # Prompt input
+    prompt = st.text_area("Enter your prompt:", placeholder="Type something...")
+
+    if st.button("Generate Response"):
+        if prompt:
+            # Generate text based on the prompt
+            with st.spinner("Generating response..."):
+                generated_text = pipe(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.")
+
+
+# Running Gradio and Streamlit interfaces
+if __name__ == "__main__":
+    st.sidebar.title("Choose an Interface")
+    interface = st.sidebar.radio("Select", ("Streamlit", "Gradio"))
+
+    if interface == "Streamlit":
+        streamlit_interface()
     else:
-        st.write("Please enter a prompt to generate a response.")
+        demo = gr.ChatInterface(
+            respond,
+            additional_inputs=[
+                gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+                gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+                gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+                gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
+            ],
+        )
+        demo.launch()