File size: 1,508 Bytes
22a519a 96255f8 22a519a 6e8709b 22a519a 9d1baf7 22a519a |
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 |
import torch
from transformers import DistilBertTokenizerFast, DistilBertForQuestionAnswering
model_name = "distilbert-base-cased"
tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
model = DistilBertForQuestionAnswering.from_pretrained(model_name)
def format_response(start_index, end_index, raw_answer):
answer_tokens = tokenizer.convert_tokens_to_string([tokenizer.convert_ids_to_tokens(i)[0] for i in range(start_index, end_index+1)])
return answer_tokens.strip()
def get_answers(question, context):
inputs = tokenizer.encode_plus(question, context, return_tensors="pt")
start_scores, end_scores = model(**inputs).values()
start_index = torch.argmax(start_scores)
end_index = torch.argmax(end_scores) + 1
formatted_answer = format_response(start_index, end_index - 1, context[start_index:end_index].tolist())
return formatted_answer
def main():
print("Hi! I am a simple AI chatbot built using Hugging Face.")
print("Type 'quit' to exit the program.")
while True:
query = input("Your Question: ")#.strip()
if query.lower() == "quit":
break
else:
if len(query) > 0:
context = "The capital of France is Paris."
try:
response = get_answers(query, context)
print(f"\nResponse: {response}\n")
except Exception as e:
print(f"\nError occurred: {str(e)}\n")
if __name__ == "__main__":
main() |