Update main.py

main.py

@@ -4,9 +4,15 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 import re
 import os
-from transformers import pipeline
+import transformers 
+import torch
 
-bart_summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
+
+model_id = "meta-llama/Meta-Llama-3-8B"
+
+pipeline = transformers.pipeline(
+    "text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16}, device_map="auto"
+)
 
 
 app = FastAPI()
@@ -107,7 +113,7 @@ async def root(item: Item):
 
         article = extract_article_content(item.url)
 
-        output = summarizer(ARTICLE, max_length=130, min_length=30, do_sample=False)
+        output = pipeline("Summarize the following news article for me:", article)
     
         return {"summary": output}