| 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.") | |
| query = "" | |
| while True: | |
| query = input("Your Question: ").strip() | |
| if query.lower() == "quit": | |
| break | |
| elif 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() |