File size: 2,080 Bytes
96b2dea
1ea6737
96b2dea
27367c2
1ea6737
64c1f09
96b2dea
 
 
cb35691
9959186
96b2dea
 
9959186
1ea6737
 
 
 
 
 
 
 
 
 
96b2dea
64c1f09
27367c2
 
1ea6737
 
 
 
 
 
96b2dea
64c1f09
27367c2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
96b2dea
27367c2
 
 
0a7206e
64c1f09
 
27367c2
1ea6737
 
 
 
 
 
96b2dea
1ea6737
96b2dea
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
from fastapi import FastAPI, Request
from pydantic import BaseModel
import transformers
import torch
from fastapi.middleware.cors import CORSMiddleware


import os
access_token_read = os.getenv("DS4")
print(access_token_read)

from huggingface_hub import login
login(token = access_token_read)

# Define the FastAPI app
app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

pipe = transformers.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 = [
        {
            "role": "system",
            "content": "You are an email summarizer. Your goal is to provide a concise summary by focusing on key points, action items, and urgency."
        },
        {
            "role": "user",
            "content": f"""
            Summarize the following email by focusing on the key points, action items, and urgency.

            Email Details:
            Subject: {subject}
            Sender: {sender}
            Recipients: {recipients}

            Body:
            {body}

            Provide a concise summary of email body in points that includes important information, if any actions are required, and the priority of the email.
            """
        }
    ]
    prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    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
    outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
    
    return {"summary": outputs}