Update app.py

app.py

@@ -1,26 +1,22 @@
 import streamlit as st
 import pandas as pd
 from transformers import pipeline
-import torch
 
 # Load translation and summarization pipelines
 translator = pipeline("translation_ru_to_en", model="Helsinki-NLP/opus-mt-ru-en")
 summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
+title_generator = pipeline("text2text-generation", model="facebook/bart-large-cnn")
 
-torch.manual_seed(0)
-generator = pipeline('text-generation', model = 'openai-community/gpt2')
-
-
-
-# Function to translate and summarize text
-def translate_and_summarize(text):
+# Function to translate, summarize, and generate title
+def translate_summarize_and_generate_title(text):
     translated_text = translator(text)[0]['translation_text']
-    summary = 'News Alert. ' + summarizer(translated_text, max_length=140, min_length=110, do_sample=False)[0]['summary_text'] 
-    return summary
+    summary = 'News Alert. ' + summarizer(translated_text, max_length=140, min_length=110, do_sample=False)[0]['summary_text']
+    title = title_generator(translated_text, max_length=10, min_length=5, do_sample=False)[0]['generated_text']
+    return summary, title
 
 # Streamlit interface
 def main():
-    st.title("CSV Translator and Summarizer")
+    st.title("CSV Translator, Summarizer, and Title Generator")
 
     # File uploader
     uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
@@ -29,19 +25,15 @@ def main():
         data = pd.read_csv(uploaded_file)
         # Check if 'Description' and 'Published' columns exist
         if 'Description' in data.columns and 'Published' in data.columns:
-            # Apply translation and summarization based on 'Published' column
-            data['Summary'] = data.apply(
-                lambda row: translate_and_summarize(row['Description']) if pd.isna(row['Published']) else "", axis=1
+            # Apply translation, summarization, and title generation based on 'Published' column
+            data[['Summary', 'Generated Title']] = data.apply(
+                lambda row: translate_summarize_and_generate_title(row['Description']) if pd.isna(row['Published']) else ("", ""), axis=1, result_type='expand'
             )
-            prompt = "Generate title for this news" + data['Summary']
-            
-            data['Title'] = generator(prompt, max_length = 30)
-
             # Display data in a table
-            st.write(data[['ID', 'Title', 'Summary', 'Title']])
+            st.write(data[['ID', 'Title', 'Generated Title', 'Summary']])
 
         else:
             st.error("Uploaded CSV does not contain required 'Description' and 'Published' columns.")
 
 if __name__ == "__main__":
-    main()
+    main()