Update app.py

app.py

@@ -2,17 +2,22 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import base64
 import pdfplumber
-from transformers import pipeline
+from transformers import DistilBertForSequenceClassification, DistilBertTokenizerFast, pipeline
 import torch
 from typing import List, Dict
 
-
 # Initialize FastAPI app
 app = FastAPI()
 
-# Load the pre-trained BERT model for contract clause classification
-classifier = pipeline("text-classification", model="distilbert-base-uncased")
+# Load the pre-trained model and tokenizer for classification
+# DistilBERT model - it's important to fine-tune this model for your task, but we'll use it as-is for now
+model_name = "distilbert-base-uncased"
+
+tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
+model = DistilBertForSequenceClassification.from_pretrained(model_name)
 
+# Use Hugging Face's pipeline for text classification
+classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)
 
 # Function to extract text from PDF file
 def extract_text_from_pdf(pdf_data: bytes) -> str: