Update app.py

app.py

@@ -1,35 +1,21 @@
-from flask import Flask, request, jsonify
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-
-app = Flask(__name__)
-
-# Load tokenizer and model once when the server starts
-tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")
-model = AutoModelForCausalLM.from_pretrained("openai/gpt-oss-20b")
-
-# Move model to GPU if available, else CPU
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
-@app.route('/generate', methods=['POST'])
-def generate_text():
-    data = request.get_json()
-    prompt = data.get('prompt')
-
-    if not prompt:
-        return jsonify({'error': 'No prompt provided'}), 400
-
-    # Tokenize input and move tensors to device
-    inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
-
-    # Generate output tokens (you can tweak max_length)
-    outputs = model.generate(inputs, max_length=50, do_sample=True, top_k=50, top_p=0.95)
-
-    # Decode tokens to string
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    return jsonify({'generated_text': generated_text})
-
-if __name__ == '__main__':
-    app.run(debug=True)
+import os
+from openai import OpenAI
+
+client = OpenAI(
+    base_url="https://router.huggingface.co/v1",
+    api_key=os.environ["HF_TOKEN"],
+)
+
+stream = client.chat.completions.create(
+    model="openai/gpt-oss-20b",
+    messages=[
+        {
+            "role": "user",
+            "content": "What is the capital of France?"
+        }
+    ],
+    stream=True,
+)
+
+for chunk in stream:
+    print(chunk.choices[0].delta.content, end="")