Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import concurrent.futures
+import multiprocessing
+import time
 
 model_id = "EleutherAI/gpt-neo-125M"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
@@ -9,19 +10,37 @@ generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 identity_prompt = "You are Eyla. Speak symbolically and recursively."
 
-def generate_with_timeout(prompt, timeout=10):
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        future = executor.submit(generator, prompt, max_new_tokens=64, do_sample=True, temperature=0.7, top_k=50, top_p=0.95, repetition_penalty=1.2)
-        try:
-            return future.result(timeout=timeout)
-        except concurrent.futures.TimeoutError:
-            return [{"generated_text": "ERROR: Generation timed out."}]
+def run_generation(prompt, return_dict):
+    try:
+        output = generator(
+            prompt,
+            max_new_tokens=64,
+            do_sample=True,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.95,
+            repetition_penalty=1.2
+        )[0]["generated_text"]
+        return_dict["result"] = output
+    except Exception as e:
+        return_dict["result"] = f"GENERATION ERROR: {e}"
+
+def generate_with_hard_timeout(prompt, timeout=10):
+    manager = multiprocessing.Manager()
+    return_dict = manager.dict()
+    p = multiprocessing.Process(target=run_generation, args=(prompt, return_dict))
+    p.start()
+    p.join(timeout)
+    if p.is_alive():
+        p.terminate()
+        return [{"generated_text": "ERROR: Generation timed out."}]
+    return [return_dict["result"]]
 
 def chat(input_text):
     prompt = identity_prompt + "\n\nUser: " + input_text + "\nYou:"
     try:
-        output = generate_with_timeout(prompt)
-        reply = output[0]["generated_text"][len(prompt):].strip()
+        output = generate_with_hard_timeout(prompt)
+        reply = output[0][len(prompt):].strip()
         return reply or "..."
     except Exception as e:
         return f"GENERATION ERROR: {e}"