Update main.py

main.py

@@ -1,6 +1,7 @@
 from fastapi import FastAPI
 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
@@ -22,23 +23,12 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Load the model and tokenizer from Hugging Face, set device to CPU
-model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"  # Replace with an appropriate model
-tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)
-model = transformers.AutoModelForCausalLM.from_pretrained(
-    model_id, 
-    # Removed device_map and low_cpu_mem_usage to avoid the need for 'accelerate'
-)
 
-# Set up the text generation pipeline for CPU
-pipeline = transformers.pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=150,
-    temperature=0.7,
-    device=-1  # Force CPU usage
-)
+
+
+
+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):
@@ -47,11 +37,38 @@ class EmailRequest(BaseModel):
     recipients: str
     body: str
 
-# Helper function to create the email prompt
+
 def create_email_prompt(subject, sender, recipients, body):
-    prompt = f"Subject: {subject}\nFrom: {sender}\nTo: {recipients}\n\n{body}\n\nSummarize this email."
+    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 that includes important information, if any actions are required, and the priority of the email.
+            """
+        }
+    ]
+    prompt = 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):
@@ -61,3 +78,4 @@ async def summarize_email(email: EmailRequest):
     summary = pipeline(prompt)[0]["generated_text"]
     
     return {"summary": summary}
+