Update app.py

app.py

@@ -1,34 +1,47 @@
-import streamlit as st
-from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+import pandas as pd
+import requests
 import json
-from datasets import load_dataset
+from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+import streamlit as st
+
+# Fetch the dataset
+url = 'https://huggingface.co/spaces/AreesaAshfaq/ContentGenerator/raw/main/article-2020.json'
+response = requests.get(url)
+data = response.json()
 
-# Load the dataset
-dataset = load_dataset('json', data_files={'data': 'https://huggingface.co/spaces/AreesaAshfaq/ContentGenerator/blob/main/article-2020.json'})['data']
+# Convert to DataFrame
+df = pd.DataFrame(data)
+#print(df.head())
 
-# Load the model and tokenizer
-tokenizer = RagTokenizer.from_pretrained('facebook/rag-sequence-nq')
-retriever = RagRetriever.from_pretrained('facebook/rag-sequence-nq')
-model = RagSequenceForGeneration.from_pretrained('facebook/rag-sequence-nq')
 
-# Define the article generation function
-def generate_article(query, dataset):
+# Initialize the tokenizer, retriever, and model
+tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
+retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq")
+model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq")
+
+def generate_article(query):
     inputs = tokenizer(query, return_tensors="pt")
-    retriever.set_ctx_data(dataset)
-    generated_ids = model.generate(
-        input_ids=inputs["input_ids"],
-        attention_mask=inputs["attention_mask"],
-        max_length=200
-    )
-    return tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-
-# Streamlit UI
-st.title('Tech Article Generator')
-
-topics = [item['title'] for item in dataset]
-selected_topic = st.selectbox('Select a topic', topics)
-
-if st.button('Generate Article'):
-    query = selected_topic
-    article = generate_article(query, dataset)
-    st.write(article)
+    input_ids = inputs["input_ids"]
+    
+    # Generate articles
+    with torch.no_grad():
+        outputs = model.generate(input_ids=input_ids, num_beams=4, max_length=200)
+    
+    article = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return article
+
+
+# Set up the Streamlit app
+st.title('Tech Article Generator using RAG')
+
+# User input for query
+query = st.text_input("Enter a topic or keyword:")
+
+if st.button("Generate Article"):
+    if query:
+        # Generate article
+        article = generate_article(query)
+        st.subheader("Generated Article:")
+        st.write(article)
+    else:
+        st.error("Please enter a query to generate an article.")