Zhannie
Zhannie

Reputation: 217

Fix ERROR:root:An error occurred: MPS does not support cumsum op with int64 input

I am working on a project trying to summarize a bunch of movie reviews using an LLM. I get that error because I'm using a mac chip m3. I don't want to use CPU because I get like 700 clusters, I want to leverage my M3. What do to? This is my code:

from domain.interfaces import SummarizationServiceInterface
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM`
import torch`
import os`

HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")`
class SummarizationService(SummarizationServiceInterface):
   """Summarizes text using a pre-trained model."""`

   def __init__(self, model_name: str = "sshleifer/distilbart-cnn-12-6", device: str = "mps"):
       """
       Initializes the summarization service.

       Args:
           model_name (str): Name of the pre-trained model to load.
           device (str): Device to run the model on ('cpu' or 'mps').
       """
       self.device = torch.device(device)
       self.tokenizer = AutoTokenizer.from_pretrained(model_name)
       self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(self.device)

   def summarize(self, texts: list[str]) -> str:
       """
       Summarizes a list of texts.

       Args:
           texts (list[str]): A list of text to summarize.

       Returns:
           str: The summarized text.
       """
       combined_text = " ".join(texts)
       inputs = self.tokenizer(
           combined_text, max_length=512, return_tensors="pt", truncation=True
       ).to(self.device)

       summary_ids = self.model.generate(
           inputs["input_ids"], max_length=150, min_length=40, num_beams=4, early_stopping=True
       )
       return self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)

Upvotes: 0

Views: 32

Answers (0)

Related Questions