Problem with answer

#57
by mohamedachilij - opened

i have a problem with the answer, its keep giving me the prompt in the answer

import torch
from PIL import Image
from pdf2image import convert_from_path
from transformers import MllamaForConditionalGeneration, MllamaProcessor, BitsAndBytesConfig
import io
import re
import gradio as gr

----------------------------------------------------------------

CONFIGURATION DU MODÈLE LOCAL

----------------------------------------------------------------

Configuration BitsAndBytes pour économiser la mémoire

quantization_config = BitsAndBytesConfig(
load_in_4bit=True, # Activer la quantification 4-bit
bnb_4bit_compute_dtype=torch.float16 # Utiliser float16 pour les calculs
)

Chemin local vers le modèle

model_path = "./models/Llama-3.2-Vision-11B"

Chargement du modèle avec configuration 4-bit

model = MllamaForConditionalGeneration.from_pretrained(
model_path,
device_map="auto", # Répartition automatique sur GPU/CPU
quantization_config=quantization_config, # Configuration BitsAndBytes
torch_dtype=torch.float16 # Type de données pour les calculs
)

Forcer l'appel de tie_weights() pour lier les poids

model.tie_weights()

Chargement du processeur

processor = MllamaProcessor.from_pretrained(model_path)

----------------------------------------------------------------

FONCTIONS POUR LA CONVERSION PDF EN IMAGE

----------------------------------------------------------------

def convert_pdf_to_image(pdf_path, dpi=300):
"""Convertir un PDF en images depuis un fichier local."""
try:
pages = convert_from_path(pdf_path, dpi=dpi)
return pages
except Exception as e:
print(f"Erreur lors de la conversion du PDF : {e}")
return []

def resize_image(img, size=(1024, 1024)):
"""Redimensionner l'image pour qu'elle tienne dans 1024x1024 pixels."""
if img.size > size:
img.thumbnail(size)
return img

----------------------------------------------------------------

TRAITEMENT DU PDF AVEC LE MODÈLE LOCAL

----------------------------------------------------------------

def process_pdf_with_model(pdf_file, prompt, max_new_tokens=100):
"""Traiter le PDF avec le modèle local."""
pages = convert_pdf_to_image(pdf_file.name) # Utiliser le nom du fichier PDF
results = []

for i, page in enumerate(pages):
    print(f"Traitement de la page {i+1}/{len(pages)}...")

    # Redimensionner l'image
    resized_image = resize_image(page)

    # Préparer l'entrée pour le modèle avec l'image et le prompt
    inputs = processor(images=resized_image, text=f"<|image|><|begin_of_text|>>{prompt}", return_tensors="pt").to(model.device)

    # Générer la réponse
    output = model.generate(**inputs, max_new_tokens=max_new_tokens)

    # Décoder la sortie
    result = processor.decode(output[0], skip_special_tokens=True)

    # Filtrer le résultat pour enlever le prompt et les éléments indésirables
    result_cleaned = re.sub(r'<OCR/>', '', result_cleaned)  # Enlever les balises <OCR/>
    result_cleaned = re.sub(r'^\s*|\s*$', '', result_cleaned)  # Enlever les espaces en début et fin
    result_cleaned = re.sub(r'\s+', ' ', result_cleaned)  # Remplacer les espaces multiples par un seul espace

    results.append({"page": i+1, "result": result_cleaned})

return results

----------------------------------------------------------------

LANCEMENT DU SCRIPT AVEC GRADIO

----------------------------------------------------------------

def gradio_interface(pdf_file, prompt, temperature):
results = process_pdf_with_model(pdf_file, prompt, max_new_tokens=50)
return [f"Page {result['page']} : {result['result']}" for result in results]

if name == "main":
# Interface Gradio
interface = gr.Interface(
fn=gradio_interface,
inputs=[
gr.File(label="Téléchargez le fichier PDF"),
gr.Textbox(label="Prompt", value="dans le cadre d'une demande de crédit nous souhaitons classifier ce document, de quel document s'agit-il? une seule proposition"),
gr.Slider(minimum=0, maximum=1, value=0.5, label="Température")
],
outputs="text",
title="Traitement de PDF avec Llama",
description="Téléchargez un fichier PDF et entrez un prompt pour classifier le document."
)

interface.launch()

Sign up or log in to comment