Create qa_gen.py

app/pyfiles/qa_gen.py

@@ -0,0 +1,31 @@
+import os
+from pathlib import Path
+import numpy as np
+import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer 
+
+current_path = os.path.dirname(os.path.abspath(__file__))
+tokenizer_path = os.path.join(current_path, "gpt_tokenizer")
+model_path = os.path.join(current_path, "gpt2_qa_model")
+tokenizer = GPT2Tokenizer.from_pretrained(tokenizer_path) # also try gpt2-medium
+model = GPT2LMHeadModel.from_pretrained(model_path)
+def generate_text(sequence, max_new_tokens):
+    ids = tokenizer.encode(f'{sequence}', return_tensors='pt')
+    input_length = ids.size(1)
+    max_length = input_length + max_new_tokens
+    final_outputs = model.generate(
+        ids,
+        do_sample=True,
+        max_length=max_length,
+        pad_token_id=model.config.eos_token_id
+    )
+    return tokenizer.decode(final_outputs[0], skip_special_tokens=True)
+
+
+def question_awnser(prompt: str):
+    result = generate_text("Question: " + prompt + "Answer: ", 35).split('Answer: ')[1]
+    try:
+        result = result.split('.')[0] + '.'
+    except Exception as e:
+        print(e)
+    return result