app/pipeline.py

import logging

import torch
from transformers import pipeline

logger = logging.getLogger(__name__)
logging.basicConfig(filename="pipeline.log", level=logging.INFO, datefmt="%Y-%m-%d %H:%M:%S")

class Pipeline:
 def __init__(self, model_name: str = "Qwen/Qwen2.5-7B-Instruct"):
 self.torch_pipe = pipeline(
 "text-generation",
 model_name,
 torch_dtype="auto",
 device_map="auto"
 )
 self.device = self._determine_device()
 logger.info(f"device type: {self.device}")
 self.messages = [
 {"role": "system", "content": """You are an expert flashcard creator.
 - You ALWAYS include a single knowledge item per flashcard.
 - You ALWAYS respond in valid JSON format.
 - You ALWAYS make flashcards accurate and comprehensivce.
 - If the text includes code snippets, you consider snippets a knowledge item testing the user's understanding of how to write the code and what it does.

 Format responses like the example below.

 EXAMPLE:
 [
 {"question": "What is AI?", "answer": "Artificial Intelligence."},
 {"question": "What is ML?", "answer": "Machine Learning."}
 ]
 """},
 ]

 def extract_flashcards(self, content: str = "", max_new_tokens: int = 1024) -> str:
 user_prompt = {"role": "user", "content": content}
 self.messages.append(user_prompt)
 try:
 response_message = self.torch_pipe(
 self.messages,
 max_new_tokens=max_new_tokens
 )[0]["generated_text"][-1]
 return response_message
 except Exception as e:
 logger.error(f"Error extracting flashcards: {str(e)}")
 raise ValueError(f"Error extraction flashcards: {str(e)}")

 def _determine_device(self) -> torch.device:
 if torch.cuda.is_available():
 return torch.device("cuda")
 elif torch.backends.mps.is_available():
 return torch.device("mps")
 else:
 return torch.device("cpu")