distilBertSquad / app.py
Dofla's picture
Update app.py
c5de0fe verified
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import torch
import gradio as gr
# Charger le tokenizer depuis Hugging Face Spaces
tokenizer = AutoTokenizer.from_pretrained("Dofla/distilbert-squad")
# Charger le modèle depuis Hugging Face Spaces
model = AutoModelForQuestionAnswering.from_pretrained("Dofla/distilbert-squad")
def answer_question(context, question):
inputs = tokenizer.encode_plus(question, context, return_tensors="pt", padding=True, truncation=True)
start_logits, end_logits = model(**inputs)
outputs = model(**inputs)
start_logits = outputs.start_logits
end_logits = outputs.end_logits
# Assurez-vous que les logits sont des tenseurs
start_index = torch.argmax(start_logits, dim=1).item()
end_index = torch.argmax(end_logits, dim=1).item() + 1
answer = tokenizer.decode(inputs["input_ids"][0][start_index:end_index])
return answer
# Créer une interface Gradio pour l'inférence
iface = gr.Interface(
fn=answer_question,
inputs=[
gr.Textbox(lines=7, label="Contexte"),
gr.Textbox(lines=1, label="Question")
],
outputs="text",
title="Question Answering with Fine-Tuned Model"
)
# Lancer l'interface
iface.launch('share=True')