Spaces:
Sleeping
Sleeping
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 |