Update app.py

app.py

@@ -1,16 +1,19 @@
 from flask import Flask, request, Response, stream_with_context
-from llama_cpp import Llama
-import json
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
+import torch
 
 app = Flask(__name__)
 
-# Model setup - Using the Gemma-3-1B Heretic GGUF
-# Ensure this file is uploaded to your Space as well
-llm = Llama(
-    model_path="gemma-3-1b-it-heretic-extreme-uncensored-abliterated.Q4_K_S.gguf",
-    n_ctx=1024,
-    n_threads=2, # Optimized for HF Free CPU
-    verbose=False
+model_id = "google/gemma-3-1b-it" # Using the official IT model
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+# Load in 4-bit to fit easily and run faster on CPU
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    low_cpu_mem_usage=True,
+    load_in_4bit=True 
 )
 
 @app.route('/generate', methods=['POST'])
@@ -18,20 +21,26 @@ def generate():
     data = request.json
     prompt = data.get("prompt", "")
     
-    # System instruction for Jarvis personality
-    full_prompt = f"<|system|>You are Jarvis, a witty and concise AI assistant. Respond briefly.<|user|>{prompt}<|assistant|>"
+    # Format for Gemma 3
+    messages = [
+        {"role": "system", "content": "You are Jarvis. Be concise."},
+        {"role": "user", "content": prompt}
+    ]
+    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
+
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    
+    # Run generation in a separate thread so we can yield tokens immediately
+    generation_kwargs = dict(input_ids=inputs, streamer=streamer, max_new_tokens=128)
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
 
     @stream_with_context
-    def generate_tokens():
-        # stream=True is the key for word-by-word
-        output = llm(full_prompt, stream=True, max_tokens=150)
-        for chunk in output:
-            token = chunk['choices'][0]['text']
-            if token:
-                # Yielding the token immediately sends it to your laptop
-                yield token
-
-    return Response(generate_tokens(), mimetype='text/plain')
+    def stream_words():
+        for new_text in streamer:
+            yield new_text
+
+    return Response(stream_words(), mimetype='text/plain')
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)