working-rabbits / app.py
kengui's picture
Update app.py
d6deb7d verified
raw
history blame
2.19 kB
import gradio as gr
from transformers import pipeline
from PyPDF2 import PdfReader
from huggingface_hub import InferenceClient
import requests
from PIL import Image
import io
pipe = pipeline("text2text-generation", model="asach/simpleT5-resume-summarization")
my_key = "YOUR_HUGGING_FACE_API_KEY"
client = InferenceClient(api_key=my_key)
def process_pdf(pdf_file):
reader = PdfReader(pdf_file.name)
text = ""
for page in reader.pages:
text += page.extract_text()
summary = pipe(text, max_length=150, min_length=30)[0]['generated_text']
agent_desc = """
You are an AI agent helps a user generate a prompt to feed into an AI image
generation model based on a summary of their resume given to you. The image should depict a rabbit
within the the career field related to the summary. Encapsulate the image prompt between
two '---' marks.
"""
messages = [
{"role": "user", "content": agent_desc},
{"role": "user", "content": summary}
]
response_text = ""
stream = client.chat.completions.create(
model='meta-llama/Llama-3.2-3B-Instruct',
messages=messages,
max_tokens=700,
stream=True
)
for chunk in stream:
response_text += chunk.choices[0].delta.content
image_prompt = response_text.split('---')[1].strip()
API_URL = "https://api-inference.huggingface.co/models/CompVis/stable-diffusion-v1-4"
headers = {"Authorization": f"Bearer {my_key}"}
def query(payload):
response = requests.post(API_URL, headers=headers, json=payload)
return response.content
image_bytes = query({"inputs": image_prompt})
image = Image.open(io.BytesIO(image_bytes))
return summary, image
pdf_input = gr.File(label="Upload PDF Resume")
summary_output = gr.Textbox(label="Resume Summary")
image_output = gr.outputs.Image(label="Generated Image")
gr.Interface(
fn=process_pdf,
inputs=pdf_input,
outputs=[summary_output, image_output],
title="Resume Summarization and Image Generation",
description="Upload your PDF resume to get a summary and a related image of a rabbit.",
allow_flagging="never"
).launch()