Update main.py

main.py

@@ -1,17 +1,16 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from pydantic import BaseModel
+import transformers
 import torch
-from transformers import pipeline
 from fastapi.middleware.cors import CORSMiddleware
-import os
-from huggingface_hub import login
 
-# Get access token from environment variable
-access_token_read = os.getenv('DS4')
+
+import os
+access_token_read = os.getenv("DS4")
 print(access_token_read)
 
-# Login to Hugging Face Hub
-login(token=access_token_read)
+from huggingface_hub import login
+login(token = access_token_read)
 
 # Define the FastAPI app
 app = FastAPI()
@@ -23,21 +22,17 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+pipe = transformers.pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.bfloat16, device_map="auto")
 
 
 
-
-pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.bfloat16, device_map="auto")
-
-
 # Define the request model for email input
 class EmailRequest(BaseModel):
     subject: str
     sender: str
     recipients: str
     body: str
-
-
+    
 def create_email_prompt(subject, sender, recipients, body):
     messages = [
         {
@@ -57,7 +52,7 @@ def create_email_prompt(subject, sender, recipients, body):
             Body:
             {body}
 
-            Your Task : "Provide a concise summary that includes important information, if any actions are required, and the priority of the email."
+            Provide a concise summary of email body in points that includes important information, if any actions are required, and the priority of the email.
             """
         }
     ]
@@ -65,17 +60,12 @@ def create_email_prompt(subject, sender, recipients, body):
     return prompt
 
 
-
-
-
-
 # Define the FastAPI endpoint for email summarization
 @app.post("/summarize-email/")
 async def summarize_email(email: EmailRequest):
     prompt = create_email_prompt(email.subject, email.sender, email.recipients, email.body)
     
     # Use the pipeline to generate the summary
-    summary = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
+    outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
     
-    return {"summary": summary}
-
+    return {"summary": outputs}