Create api.py

api.py

@@ -0,0 +1,79 @@
+from fastapi import FastAPI, UploadFile, File, Form
+from fastapi.middleware.cors import CORSMiddleware
+import torch
+from transformers import pipeline
+import PyPDF2
+from docx import Document
+import io
+
+app = FastAPI(
+    title="Text Summarization API",
+    description="API for summarizing text and documents using Falcon's text summarization model"
+)
+
+# Configure CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Initialize the summarization pipeline
+device = "cuda" if torch.cuda.is_available() else "cpu"
+summarization_pipe = pipeline("summarization", model="Falconsai/text_summarization", device=device)
+
+def extract_text_from_pdf(file_bytes):
+    pdf_reader = PyPDF2.PdfReader(io.BytesIO(file_bytes))
+    text = ""
+    for page in pdf_reader.pages:
+        text += page.extract_text()
+    return text
+
+def extract_text_from_docx(file_bytes):
+    doc = Document(io.BytesIO(file_bytes))
+    text = ""
+    for paragraph in doc.paragraphs:
+        text += paragraph.text + "\n"
+    return text
+
+@app.post("/summarize/text")
+async def summarize_text(text: str = Form(...)):
+    """
+    Summarize text input
+    """
+    if not text:
+        return {"error": "Please provide text to summarize"}
+    
+    summary = summarization_pipe(text)
+    return {"summary": summary[0]['summary_text']}
+
+@app.post("/summarize/file")
+async def summarize_file(file: UploadFile = File(...)):
+    """
+    Summarize text from a PDF or DOCX file
+    """
+    contents = await file.read()
+    file_name = file.filename.lower()
+    
+    try:
+        if file_name.endswith('.pdf'):
+            text = extract_text_from_pdf(contents)
+        elif file_name.endswith('.docx'):
+            text = extract_text_from_docx(contents)
+        else:
+            return {"error": "Unsupported file format. Please upload a PDF or DOCX file."}
+        
+        if not text:
+            return {"error": "Could not extract text from the file"}
+        
+        summary = summarization_pipe(text)
+        return {"summary": summary[0]['summary_text']}
+    
+    except Exception as e:
+        return {"error": f"Error processing file: {str(e)}"}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)