Spaces:
Runtime error
Runtime error
File size: 4,025 Bytes
2cc4144 60c61c5 2cc4144 5861c8f 2cc4144 7717dbf 2cc4144 81d062a 751e496 2cc4144 5861c8f 2cc4144 81d062a 2cc4144 81d062a 2cc4144 5861c8f 2cc4144 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
import gradio as gr
from gradio.components import Textbox
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, T5ForConditionalGeneration
from peft import PeftModel
import torch
import datasets
from sentence_transformers import SentenceTransformer, util
import math
import re
from nltk import sent_tokenize, word_tokenize
import nltk
nltk.download('punkt')
# Load bi encoder
bi_encoder = SentenceTransformer('legacy107/multi-qa-mpnet-base-dot-v1-wikipedia-search')
bi_encoder.max_seq_length = 256
top_k = 3
# Load your fine-tuned model and tokenizer
model_name = "legacy107/flan-t5-large-ia3-wiki2-100-merged"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)
max_length = 512
max_target_length = 200
# Load your dataset
dataset = datasets.load_dataset("legacy107/qa_wikipedia_retrieved_chunks", split="test")
dataset = dataset.shuffle()
dataset = dataset.select(range(10))
# Context chunking
def chunk_splitter(context, chunk_size=100, overlap=0.20):
overlap_size = chunk_size * overlap
sentences = nltk.sent_tokenize(context)
chunks = []
text = sentences[0]
if len(sentences) == 1:
chunks.append(text)
i = 1
while i < len(sentences):
text += " " + sentences[i]
i += 1
while i < len(sentences) and len(nltk.word_tokenize(f"{text} {sentences[i]}")) <= chunk_size:
text += " " + sentences[i]
i += 1
text = text.replace('\"','"').replace("\'","'").replace('\n\n\n'," ").replace('\n\n'," ").replace('\n'," ")
chunks.append(text)
if (i >= len(sentences)):
break
j = i - 1
text = sentences[j]
while j >= 0 and len(nltk.word_tokenize(f"{sentences[j]} {text}")) <= overlap_size:
text = sentences[j] + " " + text
j -= 1
return chunks
def retrieve_context(query, contexts):
corpus_embeddings = bi_encoder.encode(contexts, convert_to_tensor=True, show_progress_bar=False)
question_embedding = bi_encoder.encode(query, convert_to_tensor=True, show_progress_bar=False)
question_embedding = question_embedding
hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)
hits = hits[0]
hits = sorted(hits, key=lambda x: x['score'], reverse=True)
return " ".join([contexts[hit['corpus_id']] for hit in hits[0:top_k]]).replace("\n", " ")
# Define your function to generate answers
def generate_answer(question, context, title, ground):
contexts = chunk_splitter(context)
context = retrieve_context(question, contexts)
# Combine question and context
input_text = f"question: {question} context: {context}"
# Tokenize the input text
input_ids = tokenizer(
input_text,
return_tensors="pt",
padding="max_length",
truncation=True,
max_length=max_length,
).input_ids
# Generate the answer
with torch.no_grad():
generated_ids = model.generate(input_ids=input_ids, max_new_tokens=max_target_length)
# Decode and return the generated answer
generated_answer = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
return generated_answer, context
# Define a function to list examples from the dataset
def list_examples():
examples = []
for example in dataset:
context = example["article"]
question = example["question"]
answer = example["answer"]
title = example["title"]
examples.append([question, context, title, answer])
return examples
# Create a Gradio interface
iface = gr.Interface(
fn=generate_answer,
inputs=[
Textbox(label="Question"),
Textbox(label="Context"),
Textbox(label="Article title"),
Textbox(label="Ground truth")
],
outputs=[
Textbox(label="Generated Answer"),
Textbox(label="Retrieved Context")
],
examples=list_examples()
)
# Launch the Gradio interface
iface.launch() |