summarizer

summarizer
git clone https://git.prime8.dev/summarizer.git
Log | Files | Refs | README | LICENSE

main.py (2512B)


      1 import os
      2 import pickle
      3 from uuid import uuid4 as uuid
      4 
      5 from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
      6 from fastapi import FastAPI, HTTPException, UploadFile
      7 from fastapi.responses import HTMLResponse
      8 
      9 address = os.environ.get("KAFKA_ADDRESS", "localhost:9092")
     10 
     11 
     12 app = FastAPI()
     13 
     14 
     15 @app.get("/", response_class=HTMLResponse)
     16 async def index():
     17     return HTMLResponse(
     18         content="""
     19     <html>
     20         <form enctype="multipart/form-data" method="post">
     21             <input type="file" name="file" accept="audio/*" />
     22             <button type="submit">Submit</button>
     23         </form>
     24     </html>
     25     """
     26     )
     27 
     28 
     29 @app.post("/")
     30 async def summarize(file: UploadFile):
     31     job_id = str.encode(str(uuid()))
     32 
     33     consumer = None
     34     producer = None
     35 
     36     try:
     37         consumer = AIOKafkaConsumer(
     38             "transcription.output",
     39             bootstrap_servers=address,
     40             consumer_timeout_ms=300000,
     41         )
     42         await consumer.start()
     43 
     44         producer = AIOKafkaProducer(
     45             max_request_size=26214400, bootstrap_servers=address
     46         )
     47         await producer.start()
     48 
     49         print("transcribing file")
     50         await producer.send_and_wait(
     51             "transcription.input",
     52             key=job_id,
     53             value=pickle.dumps(
     54                 {
     55                     "name": file.filename,
     56                     "data": await file.read(),
     57                 }
     58             ),
     59         )
     60 
     61         text = ""
     62         async for message in consumer:
     63             if message.key == job_id:
     64                 text = message.value.decode()
     65                 break
     66 
     67         print("summarizing text")
     68         await consumer.stop()
     69         consumer = AIOKafkaConsumer(
     70             "summarization.output",
     71             bootstrap_servers=address,
     72             consumer_timeout_ms=300000,
     73         )
     74         await consumer.start()
     75 
     76         await producer.send_and_wait(
     77             "summarization.input", key=job_id, value=str.encode(text)
     78         )
     79 
     80         summary = ""
     81         async for message in consumer:
     82             if message.key == job_id:
     83                 summary = message.value.decode()
     84                 break
     85 
     86         print("request completed")
     87         return {
     88             "text": text,
     89             "summary": summary,
     90         }
     91 
     92     except Exception as e:
     93         print(f"Error processing request: {e}")
     94         raise HTTPException(status_code=500, detail=str(e))
     95 
     96     finally:
     97         if consumer:
     98             await consumer.stop()
     99         if producer:
    100             await producer.stop()