The_AI

app.py

@@ -1,17 +1,28 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import T5ForConditionalGeneration, T5Tokenizer, pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import nltk
 from youtube_transcript_api import YouTubeTranscriptApi
 
 # Download NLTK data
 nltk.download('punkt')
 
-# Initialize the image captioning pipeline
+# Load models and tokenizers
+summary_model_name = 'utrobinmv/t5_summary_en_ru_zh_base_2048'
+summary_model = T5ForConditionalGeneration.from_pretrained(summary_model_name)
+summary_tokenizer = T5Tokenizer.from_pretrained(summary_model_name)
+
+tag_tokenizer = AutoTokenizer.from_pretrained("fabiochiu/t5-base-tag-generation")
+tag_model = AutoModelForSeq2SeqLM.from_pretrained("fabiochiu/t5-base-tag-generation")
+
 captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-# Load the tokenizer and model for tag generation
-tokenizer = AutoTokenizer.from_pretrained("fabiochiu/t5-base-tag-generation")
-model = AutoModelForSeq2SeqLM.from_pretrained("fabiochiu/t5-base-tag-generation")
+# Function to summarize text
+def summarize_text(text, prefix):
+    src_text = prefix + text
+    input_ids = summary_tokenizer(src_text, return_tensors="pt")
+    generated_tokens = summary_model.generate(**input_ids)
+    result = summary_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+    return result[0]
 
 # Function to fetch YouTube transcript
 def fetch_transcript(url):
@@ -27,32 +38,38 @@ def fetch_transcript(url):
 st.title("Multi-purpose Machine Learning App: WAVE_AI")
 
 # Create tabs for different functionalities
-tab1, tab2, tab3 = st.tabs(["Text Tag Generation", "Image Captioning", "YouTube Transcript"])
+tab1, tab2, tab3, tab4 = st.tabs(["Text Summarization", "Text Tag Generation", "Image Captioning", "YouTube Transcript"])
 
-# Image Captioning Tab
+# Text Summarization Tab
 with tab1:
+    st.header("Text Summarization")
+
+    input_text = st.text_area("Enter the text to summarize:", height=300)
+
+    if st.button("Generate Summaries"):
+        if input_text:
+            title1 = summarize_text(input_text, 'summary: ')
+            title2 = summarize_text(input_text, 'summary brief: ')
+            st.write("### Title 1")
+            st.write(title1)
+            st.write("### Title 2")
+            st.write(title2)
+        else:
+            st.warning("Please enter some text to summarize.")
+
+# Text Tag Generation Tab
+with tab2:
     st.header("Text Tag Generation")
     
-    # Text area for user input
     text = st.text_area("Enter the text for tag extraction:", height=200)
     
-    # Button to generate tags
     if st.button("Generate Tags"):
         if text:
             try:
-                # Tokenize and encode the input text
-                inputs = tokenizer([text], max_length=512, truncation=True, return_tensors="pt")
-                
-                # Generate tags
-                output = model.generate(**inputs, num_beams=8, do_sample=True, min_length=10, max_length=64)
-                
-                # Decode the output
-                decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
-                
-                # Extract unique tags
+                inputs = tag_tokenizer([text], max_length=512, truncation=True, return_tensors="pt")
+                output = tag_model.generate(**inputs, num_beams=8, do_sample=True, min_length=10, max_length=64)
+                decoded_output = tag_tokenizer.batch_decode(output, skip_special_tokens=True)[0]
                 tags = list(set(decoded_output.strip().split(", ")))
-                
-                # Display the tags
                 st.write("**Generated Tags:**")
                 st.write(tags)
             except Exception as e:
@@ -60,36 +77,27 @@ with tab1:
         else:
             st.warning("Please enter some text to generate tags.")
 
-# Text Tag Generation Tab
-with tab2:
+# Image Captioning Tab
+with tab3:
     st.header("Image Captioning Extractor")
     
-    # Input for image URL
     image_url = st.text_input("Enter the URL of the image:")
     
-    # If an image URL is provided
     if image_url:
         try:
-            # Display the image
             st.image(image_url, caption="Provided Image", use_column_width=True)
-            
-            # Generate the caption
             caption = captioner(image_url)
-            
-            # Display the caption
             st.write("**Generated Caption:**")
             st.write(caption[0]['generated_text'])
         except Exception as e:
             st.error(f"An error occurred: {e}")
 
 # YouTube Transcript Tab
-with tab3:
+with tab4:
     st.header("YouTube Video Transcript Extractor")
     
-    # Input for YouTube URL
     youtube_url = st.text_input("Enter YouTube URL:")
     
-    # Button to get transcript
     if st.button("Get Transcript"):
         if youtube_url:
             transcript = fetch_transcript(youtube_url)