dodano accelerate i lepsze logowanie błędów

app.py

@@ -1,31 +1,62 @@
+import os
 import spaces
 import gradio as gr
 import transformers
 import torch
-import os
 from huggingface_hub import login
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 model_id = "meta-llama/Meta-Llama-3.1-8B"
 
 @spaces.GPU(duration=60)
 def load_pipeline():
-    # Zaloguj się używając tokena
-    login(token=os.environ.get("MY_API_LLAMA_3_1"))
+    try:
+        # Zaloguj się używając tokena
+        login(token=os.environ.get("MY_API_LLAMA_3_1"))
+        logger.info("Login successful")
 
-    return transformers.pipeline(
-        "text-generation",
-        model=model_id,
-        model_kwargs={"torch_dtype": torch.bfloat16},
-        device_map="auto"
-    )
+        if torch.cuda.is_available():
+            logger.info(f"GPU available: {torch.cuda.get_device_name(0)}")
+            device_map = "auto"
+            torch_dtype = torch.bfloat16
+        else:
+            logger.warning("No GPU available, using CPU")
+            device_map = "cpu"
+            torch_dtype = torch.float32
+
+        pipeline = transformers.pipeline(
+            "text-generation",
+            model=model_id,
+            model_kwargs={"torch_dtype": torch_dtype},
+            device_map=device_map
+        )
+        logger.info("Model loaded successfully")
+        return pipeline
+    except Exception as e:
+        logger.error(f"Error loading model: {str(e)}")
+        raise
 
-pipeline = load_pipeline()
+try:
+    pipeline = load_pipeline()
+except Exception as e:
+    logger.error(f"Failed to load pipeline: {str(e)}")
+    pipeline = None
 
 def generate_response(chat, kwargs):
-    output = pipeline(chat, **kwargs)[0]['generated_text']
-    if output.endswith("</s>"):
-        output = output[:-4]
-    return output
+    if pipeline is None:
+        return "Model nie został załadowany poprawnie. Proszę spróbować później."
+
+    try:
+        output = pipeline(chat, **kwargs)[0]['generated_text']
+        if output.endswith("</s>"):
+            output = output[:-4]
+        return output
+    except Exception as e:
+        logger.error(f"Error generating response: {str(e)}")
+        return f"Wystąpił błąd podczas generowania odpowiedzi: {str(e)}"
 
 def function(prompt, history=[]):
     chat = "<s>"
@@ -42,11 +73,7 @@ def function(prompt, history=[]):
         seed=1337
     )
 
-    try:
-        output = generate_response(chat, kwargs)
-        return output
-    except:
-        return ''
+    return generate_response(chat, kwargs)
 
 # Interfejs Gradio
 interface = gr.ChatInterface(

requirements.txt

@@ -4,3 +4,4 @@ numpy<2
 torch
 transformers
 bitsandbytes
+accelerate