Update app.py

app.py

@@ -1,47 +1,22 @@
-import gradio as gr
+from fastapi import FastAPI
+from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
-# Initialize model and tokenizer
-model_name = "deepseek-ai/deepseek-llm-1.3b"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+app = FastAPI()
 
-def summarize_text(text, max_length=150, min_length=50):
-    """
-    Summarize the input text using the DeepSeek model
-    """
-    # Prepare the input
-    inputs = tokenizer(text, return_tensors="pt", max_length=1024, truncation=True)
-    
-    # Generate summary
-    summary_ids = model.generate(
-        inputs["input_ids"],
-        max_length=max_length,
-        min_length=min_length,
-        length_penalty=2.0,
-        num_beams=4,
-        early_stopping=True
-    )
-    
-    # Decode and return the summary
-    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-    return summary
+# Load Model & Tokenizer
+MODEL_NAME = "facebook/bart-large-cnn"  # Small & fast summarization model
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to("cpu")  # Use "cuda" if you have a GPU
 
-# Create the Gradio interface
-iface = gr.Interface(
-    fn=summarize_text,
-    inputs=[
-        gr.Textbox(label="Input Text", placeholder="Enter the text you want to summarize...", lines=10),
-        gr.Slider(minimum=50, maximum=300, value=150, label="Maximum Summary Length"),
-        gr.Slider(minimum=30, maximum=150, value=50, label="Minimum Summary Length")
-    ],
-    outputs=gr.Textbox(label="Summary"),
-    title="Text Summarization with DeepSeek",
-    description="Enter your text and get an AI-generated summary using the DeepSeek model.",
-    examples=[
-        ["The artificial intelligence revolution has transformed various sectors of the economy, from healthcare to finance. Machine learning algorithms are now capable of detecting diseases, predicting market trends, and automating complex tasks. This technological advancement has raised both excitement about the potential benefits and concerns about job displacement and ethical implications."],
-    ]
-)
+# Define input format
+class InputText(BaseModel):
+    text: str
 
-iface.launch()
+@app.post("/summarize")
+async def summarize_text(input_text: InputText):
+    inputs = tokenizer(input_text.text, return_tensors="pt", max_length=1024, truncation=True)
+    summary_ids = model.generate(inputs.input_ids, max_length=150, min_length=50, length_penalty=2.0)
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return {"summary": summary}