Spaces:
Running
Running
from flask import Flask, request, jsonify | |
from transformers import GPT2LMHeadModel, GPT2Tokenizer | |
app = Flask(__name__) | |
# Load the GPT-2 model and tokenizer | |
model_name = "gpt2" | |
model = GPT2LMHeadModel.from_pretrained(model_name) | |
tokenizer = GPT2Tokenizer.from_pretrained(model_name) | |
def generate_text(): | |
data = request.get_json() | |
input_text = data['input_text'] | |
# Tokenize the input text | |
input_ids = tokenizer.encode(input_text, return_tensors='pt') | |
# Generate text using the GPT-2 model | |
output = model.generate(input_ids, max_length=100, num_return_sequences=1) | |
generated_text = tokenizer.decode(output[0], skip_special_tokens=True) | |
return jsonify({'generated_text': generated_text}) | |
if __name__ == '__main__': | |
app.run() |