insaabbas commited on
Commit
7b16f72
·
verified ·
1 Parent(s): f24e6f8

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +22 -0
app.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from flask import Flask, request, jsonify
2
+ from flask_cors import CORS
3
+ from transformers import pipeline
4
+ import torch
5
+
6
+ app = Flask(__name__)
7
+ CORS(app)
8
+
9
+ # Load YOUR specific model directly into the Space
10
+ pipe = pipeline("text-generation", model="insaabbas/phi2-4-epoch-humor-model", device_map="auto")
11
+
12
+ @app.route("/generate", methods=["POST"])
13
+ def generate():
14
+ data = request.json
15
+ prompt = data.get("prompt", "")
16
+ # Format for Phi-2
17
+ full_prompt = f"Instruct: {prompt}\nOutput:"
18
+ result = pipe(full_prompt, max_new_tokens=100, do_sample=True, temperature=0.7)
19
+ return jsonify({"joke": result[0]['generated_text'].replace(full_prompt, "").strip()})
20
+
21
+ if __name__ == "__main__":
22
+ app.run(host="0.0.0.0", port=7860)