Update main.py

main.py

@@ -4,23 +4,24 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 import re
 import os
-from transformers import AutoModelWithLMHead, AutoTokenizer
+import transformers
+import torch
 
 app = FastAPI()
 
+model_id = "meta-llama/Meta-Llama-3-8B"
 
+try:
+    pipeline = transformers.pipeline(
+    "text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16}, device_map="auto")
 
-tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-summarize-news", use_fast=False)
-model = AutoModelWithLMHead.from_pretrained("mrm8488/t5-base-finetuned-summarize-news", use_fast=False)
+    print(pipeline("Hey how are you doing today?"))
 
-def summarize(text, max_length=150):
-  input_ids = tokenizer.encode(text, return_tensors="pt", add_special_tokens=True)
+except:
+    print("ERROR")
 
-  generated_ids = model.generate(input_ids=input_ids, num_beams=2, max_length=max_length,  repetition_penalty=2.5, length_penalty=1.0, early_stopping=True)
 
-  preds = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g in generated_ids]
 
-  return preds[0]
 
 
 @app.get("/")