Update app.py

app.py

@@ -1,43 +1,28 @@
-import os
-os.system("pip install flask")
+from fastapi import FastAPI, Query
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-from flask import Flask, request, send_file, abort
-from io import BytesIO
-from PIL import Image
-import random
+app = FastAPI()
 
-app = Flask(__name__)
+# モデルとトークナイザーのロード
+def load_prompter():
+    prompter_model = AutoModelForCausalLM.from_pretrained("microsoft/Promptist")
+    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "left"
+    return prompter_model, tokenizer
 
-# モデルのロードをシミュレートするためのスタブ
-def load_model(model_name):
-    # モデルをロードするロジックをここに追加
-    # ここでは、PILで空の画像を生成する簡易的な例を示します
-    def model(prompt):
-        img = Image.new('RGB', (256, 256), color = (73, 109, 137))
-        return img
-    return model
+prompter_model, prompter_tokenizer = load_prompter()
 
-# モデルのロード
-models_load = {f'model_{i}': load_model(f'model_{i}') for i in range(1, 7)}
+@app.get("/generate")
+async def generate(text: str = Query(..., description="Input text to be processed by the model")):
+    input_ids = prompter_tokenizer(text.strip() + " Rephrase:", return_tensors="pt").input_ids
+    eos_id = prompter_tokenizer.eos_token_id
+    outputs = prompter_model.generate(input_ids, do_sample=False, max_new_tokens=75, num_beams=8, num_return_sequences=1, eos_token_id=eos_id, pad_token_id=eos_id, length_penalty=-1.0)
+    output_texts = prompter_tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    res = output_texts[0].replace(text + " Rephrase:", "").strip()
+    return {"result": res}
 
-@app.route('/generate_image', methods=['GET'])
-def generate_image():
-    prompt = request.args.get('prompt')
-    model_str = request.args.get('model')
-    
-    if model_str not in models_load:
-        abort(404, description="Model not found")
-    
-    model = models_load[model_str]
-    noise = str(random.randint(0, 99999999999))
-    image = model(f'{prompt} {noise}')
-    
-    # 画像をバイナリで保存し、レスポンスとして返す
-    img_bytes = BytesIO()
-    image.save(img_bytes, format='PNG')
-    img_bytes.seek(0)
-    
-    return send_file(img_bytes, mimetype='image/png')
-
-if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=5000)
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)