Update app.py

app.py

@@ -15,10 +15,11 @@ translation_models = {
     'Spanish': "Helsinki-NLP/opus-mt-en-es"
 }
 
-# Initialize summarization pipeline with a specified model
-summarization_model = "facebook/bart-large-cnn"
-summarizer = pipeline("summarization", model=summarization_model)  # Removed device parameter
-tokenizer = AutoTokenizer.from_pretrained(summarization_model)
+# Initialize summarization pipelines with specified models
+summarization_models = {
+    'Scientific': "allenai/scibert_scivocab_cased",
+    'Literature': "t5-small"
+}
 
 # Initialize translation pipeline
 def get_translator(language):
@@ -45,12 +46,12 @@ def split_text(text, max_tokens=1024):
     return chunks
 
 # Helper function to summarize text
-def summarize_text(text):
+def summarize_text(text, model):
     chunks = split_text(text)
     summaries = []
     for chunk in chunks:
         try:
-            summary = summarizer(chunk, max_length=150, min_length=40, do_sample=False)[0]['summary_text']
+            summary = summarizers[model](chunk, max_length=150, min_length=40, do_sample=False)[0]['summary_text']
             summaries.append(summary)
         except Exception as e:
             print(f"Error summarizing chunk: {chunk}\nError: {e}")
@@ -68,9 +69,9 @@ def translate_text(text, language):
             return text
     return text
 
-def process_text(input_text, language):
+def process_text(input_text, model, language):
     print(f"Input text: {input_text[:500]}...")  # Show only the first 500 characters for brevity
-    summary = summarize_text(input_text)
+    summary = summarize_text(input_text, model)
     print(f"Summary: {summary[:500]}...")  # Show only the first 500 characters for brevity
     bullet_points = generate_bullet_points(summary)
     print(f"Bullet Points: {bullet_points}")
@@ -95,6 +96,7 @@ iface = gr.Interface(
     fn=process_text,
     inputs=[
         gr.Textbox(label="Input Text", placeholder="Paste your text here...", lines=10),
+        gr.Radio(choices=["Scientific", "Literature"], label="Summarization Model"),
         gr.Dropdown(choices=["Vietnamese", "Japanese", "Thai", "Spanish"], label="Translate to", value="Vietnamese")
     ],
     outputs=[
@@ -102,7 +104,7 @@ iface = gr.Interface(
         gr.Textbox(label="Translated Bullet Points", lines=10)
     ],
     title="Text to Bullet Points and Translation",
-    description="Paste any text, and the program will summarize it into bullet points. Optionally, translate the bullet points into Vietnamese, Japanese, Thai, or Spanish."
+    description="Paste any text, choose the summarization model, and optionally translate the bullet points into Vietnamese, Japanese, Thai, or Spanish."
 )
 
 iface.launch()
@@ -116,3 +118,4 @@ iface.launch()
 
 
 
+