Spaces:
Sleeping
Sleeping
File size: 2,194 Bytes
b77a775 6273efa b77a775 6273efa |
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 |
import streamlit as st
from datasets import load_dataset
from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
# Load a multilingual dataset (xnli or tydi_qa)
def load_data():
try:
# Load the 'xnli' dataset, validation split
dataset = load_dataset("xnli", split="validation")
st.write(f"Loaded {len(dataset)} examples from the 'validation' split.")
return dataset
except Exception as e:
st.write(f"Error loading 'xnli' dataset: {e}")
return None
# Initialize RAG model components
def initialize_rag():
try:
# Initialize tokenizer and retriever
tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="compressed", passages_path="./path_to_data")
model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq")
return tokenizer, retriever, model
except Exception as e:
st.write(f"Error initializing RAG components: {e}")
return None, None, None
# Main function to run the app
def main():
st.title("Multilingual RAG Translator/Answer Bot")
# Load the dataset
dataset = load_data()
if dataset is None:
st.write("Dataset could not be loaded.")
return
# Initialize RAG model components
tokenizer, retriever, model = initialize_rag()
if tokenizer is None or retriever is None or model is None:
st.write("RAG components could not be initialized.")
return
# UI to input a query
query = st.text_input("Enter your question in Urdu, Hindi, or French:")
if query:
# Tokenize the input query
inputs = tokenizer(query, return_tensors="pt")
# Retrieve relevant documents
retrieved_docs = retriever.retrieve(query)
# Generate an answer using the model
generated = model.generate(input_ids=inputs['input_ids'], context_input_ids=retrieved_docs['input_ids'])
answer = tokenizer.decode(generated[0], skip_special_tokens=True)
st.write("Answer:", answer)
# Run the Streamlit app
if __name__ == "__main__":
main()
|