OrbIQ / QA.py
Debabrath's picture
Upload QA.py
c18f457 verified
# src/qa_citation.py — Question Answering with citation
from transformers import pipeline
# Load a QA pipeline with a pre-trained model
qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
def answer_question(question, context):
"""Answer the question based on context and provide citation info."""
result = qa_pipeline(question=question, context=context)
answer = result['answer']
score = result['score']
# For citation, we’ll just return the context snippet here (could be URL in production)
citation = "Context snippet used as citation"
return {
"answer": answer,
"score": score,
"citation": citation
}
if __name__ == "__main__":
context = (
"The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in Paris, France. "
"It was built between 1887 and 1889 as the entrance to the 1889 World's Fair."
)
question = "When was the Eiffel Tower built?"
result = answer_question(question, context)
print("Answer:", result['answer'])
print("Score:", result['score'])
print("Citation:", result['citation'])