Creatingdataset / app.py
Yoxas's picture
Update app.py
6b8dbb2 verified
import gradio as gr
import pandas as pd
import spaces
from transformers import RagTokenizer, RagRetriever, RagTokenForGeneration
# Load the tokenizer and retriever
tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", use_dummy_dataset=True)
# Load the model
model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever)
# Tokenize the contexts and responses
inputs = tokenizer(contexts, return_tensors='pt', padding=True, truncation=True)
labels = tokenizer(responses, return_tensors='pt', padding=True, truncation=True)
# Extract the abstracts
abstracts = df['Abstract'].dropna().tolist()
# Load your dataset
df = pd.read_csv('10kstats.csv')
# Generate context-response pairs (abstract-question pairs)
# Here we use the abstracts as contexts and simulate questions
contexts = abstracts
responses = ["Can you tell me more about this research?" for _ in abstracts]
@spaces.GPU
def generate_response(input_text):
input_ids = tokenizer([input_text], return_tensors='pt')['input_ids']
outputs = model.generate(input_ids)
response = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
return response
# Create the Gradio interface
iface = gr.Interface(
fn=generate_response,
inputs="text",
outputs="text",
title="RAG Chatbot",
description="A chatbot powered by Retrieval-Augmented Generation (RAG) model."
)
# Launch the interface
iface.launch()