lepsze logowanie błędów, timeout zwiększony do 120s

app.py

@@ -5,6 +5,10 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 from threading import Thread
 from queue import Queue, Empty
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 model_id = "meta-llama/Meta-Llama-3.1-8B"
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ.get("MY_API_LLAMA_3_1"))
@@ -14,45 +18,61 @@ model_load_queue = Queue()
 
 def load_model():
     global model
-    if model is None:
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            token=os.environ.get("MY_API_LLAMA_3_1"),
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-            low_cpu_mem_usage=True
-        )
-    model_load_queue.put(model)
+    try:
+        if model is None:
+            logger.info("Loading model...")
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                token=os.environ.get("MY_API_LLAMA_3_1"),
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+                low_cpu_mem_usage=True
+            )
+            logger.info("Model loaded successfully")
+        model_load_queue.put(model)
+    except Exception as e:
+        logger.error(f"Error loading model: {str(e)}")
+        model_load_queue.put(None)
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=120)
 def generate_response(chat, kwargs):
     global model
-    if model is None:
-        Thread(target=load_model).start()
-        model = model_load_queue.get()
+    try:
+        if model is None:
+            logger.info("Starting model loading thread")
+            Thread(target=load_model).start()
+            model = model_load_queue.get(timeout=120)
+            if model is None:
+                return "Nie udało się załadować modelu. Proszę spróbować ponownie później."
 
-    inputs = tokenizer(chat, return_tensors="pt").to(model.device)
-    streamer = TextIteratorStreamer(tokenizer, timeout=60., skip_prompt=True, skip_special_tokens=True)
+        logger.info("Preparing input for generation")
+        inputs = tokenizer(chat, return_tensors="pt").to(model.device)
+        streamer = TextIteratorStreamer(tokenizer, timeout=120., skip_prompt=True, skip_special_tokens=True)
 
-    # Usuwamy 'seed' z kwargs, ponieważ nie jest obsługiwany przez model
-    if 'seed' in kwargs:
-        del kwargs['seed']
+        if 'seed' in kwargs:
+            del kwargs['seed']
 
-    generation_kwargs = dict(inputs, streamer=streamer, **kwargs)
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
+        generation_kwargs = dict(inputs, streamer=streamer, **kwargs)
 
-    output = ""
-    try:
-        for new_text in streamer:
-            output += new_text
-            if output.endswith("</s>"):
-                output = output[:-4]
-                break
-    except Empty:
-        print("Timeout occurred during generation")
+        logger.info("Starting generation thread")
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
 
-    return output
+        output = ""
+        try:
+            for new_text in streamer:
+                output += new_text
+                if output.endswith("</s>"):
+                    output = output[:-4]
+                    break
+        except Empty:
+            logger.warning("Timeout occurred during generation")
+
+        logger.info("Generation completed")
+        return output
+    except Exception as e:
+        logger.error(f"Error in generate_response: {str(e)}")
+        return f"Wystąpił błąd: {str(e)}"
 
 def function(prompt, history=[]):
     chat = "<s>"
@@ -67,12 +87,7 @@ def function(prompt, history=[]):
         repetition_penalty=1.0
     )
 
-    try:
-        output = generate_response(chat, kwargs)
-        return output
-    except Exception as e:
-        print(f"Error: {str(e)}")
-        return 'Wystąpił błąd podczas generowania odpowiedzi.'
+    return generate_response(chat, kwargs)
 
 interface = gr.ChatInterface(
     fn=function,