from fastapi import FastAPI
from pydantic import BaseModel
import transformers
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')
print(access_token_read)

# Login to Hugging Face Hub
login(token=access_token_read)

# Define the FastAPI app
app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    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
)

# Define the request model for email input
class EmailRequest(BaseModel):
    subject: str
    sender: str
    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."
    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 = pipeline(prompt)[0]["generated_text"]
    
    return {"summary": summary}