Update app.py

app.py

@@ -1,27 +1,20 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoFeatureExtractor
-import torch
-from PIL import Image
-import requests
+from transformers import BartTokenizer, BartForConditionalGeneration
 
-# Load the tokenizer, model, and feature extractor
-model_name = "Salesforce/BLIP-image-captioning-base"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+model_name = "facebook/bart-large-cnn"  # Example BART model
+tokenizer = BartTokenizer.from_pretrained(model_name)
+model = BartForConditionalGeneration.from_pretrained(model_name)
 
-def generate_caption(image):
-    inputs = feature_extractor(images=image, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=128, num_beams=4, return_dict_in_generate=True)
-    caption = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
-    return caption
+def generate_text(prompt):
+    inputs = tokenizer.encode("summarize: " + prompt, return_tensors="pt", max_length=1024, truncation=True)
+    summary_ids = model.generate(inputs, max_length=150, min_length=40, length_penalty=2.0, num_beams=4, early_stopping=True)
+    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
 
-# Create the Gradio interface
-interface = gr.Interface(fn=generate_caption,
-                         inputs=gr.inputs.Image(type="pil"),
+interface = gr.Interface(fn=generate_text,
+                         inputs=gr.inputs.Textbox(lines=5, placeholder="Enter Text Here..."),
                          outputs="text",
-                         title="Image Captioning with BLIP",
-                         description="Upload an image to generate a caption.")
+                         title="Text Generation with BART",
+                         description="Enter text to generate a summary.")
 
 if __name__ == "__main__":
     interface.launch()