Update server.py

server.py

@@ -1,52 +1,65 @@
 import torch
 import re
 from html import unescape
-from transformers import GPT2LMHeadModel, GPT2Tokenizer, BitsAndBytesConfig
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
 from peft import PeftModel
 from transformers import StoppingCriteria, StoppingCriteriaList
 from difflib import SequenceMatcher
 from flask import Flask, request, jsonify
 
-# Step 2: Load tokenizer
+# --------------------------
+# Step 1: Nastavení zařízení
+# --------------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"🚀 Running on device: {device}")
+
+# --------------------------
+# Step 2: Načtení tokenizeru
+# --------------------------
 model_path = "./"
 try:
     tokenizer = GPT2Tokenizer.from_pretrained(model_path)
     tokenizer.pad_token = tokenizer.eos_token
-    print("Tokenizer loaded successfully")
+    print("✅ Tokenizer loaded successfully")
 except Exception as e:
-    print(f"Error loading tokenizer: {e}")
+    print(f"❌ Error loading tokenizer: {e}")
     exit()
 
-# Step 3: Load model (kvantizace + fallback)
-quant_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16
-)
+# --------------------------
+# Step 3: Načtení modelu s fallbackem
+# --------------------------
+quant_config = None
+if torch.cuda.is_available():
+    try:
+        from transformers import BitsAndBytesConfig
+        quant_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16
+        )
+        print("✅ Using 4-bit quantization (GPU mode)")
+    except Exception as e:
+        print("⚠️ BitsAndBytes not available, continuing without quantization:", e)
+else:
+    print("💡 CPU mode — quantization disabled")
 
 try:
     base_model = GPT2LMHeadModel.from_pretrained(
         model_path,
         quantization_config=quant_config,
-        device_map={"": 0},
-        low_cpu_mem_usage=True
-    )
-    print("Base model loaded successfully (4bit quantized)")
+        device_map={"": 0} if torch.cuda.is_available() else None,
+        low_cpu_mem_usage=True,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+    ).to(device)
+    print("✅ Base model loaded successfully")
 except Exception as e:
-    print(f"Error loading base model: {e}")
-    try:
-        base_model = GPT2LMHeadModel.from_pretrained(
-            model_path,
-            low_cpu_mem_usage=True,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-        ).to("cuda:0" if torch.cuda.is_available() else "cpu")
-        print("Base model loaded without quantization")
-    except Exception as e:
-        print(f"Error loading base model without quantization: {e}")
-        exit()
+    print(f"❌ Error loading base model: {e}")
+    exit()
 
-# Step 4: Load PEFT (LoRA)
+# --------------------------
+# Step 4: Načtení PEFT (LoRA)
+# --------------------------
 try:
     model = PeftModel.from_pretrained(
         base_model,
@@ -54,16 +67,21 @@ try:
         is_trainable=False,
         device_map={"": 0} if torch.cuda.is_available() else None
     )
-    print("PEFT model loaded successfully")
+    model.to(device)
+    print("✅ PEFT model loaded successfully")
 except Exception as e:
-    print(f"Error loading PEFT model: {e}")
-    exit()
+    print(f"⚠️ Warning: Failed to load PEFT adapter, using base model. ({e})")
+    model = base_model
 
+# --------------------------
 # Step 5: System prompt
+# --------------------------
 system_prompt = """You are TinyGPT, a friendly AI assistant made by LuxAI.
 You must answer very short."""
 
+# --------------------------
 # Step 6: Stopping criteria
+# --------------------------
 class CustomStoppingCriteria(StoppingCriteria):
     def __init__(self, stop_token_id):
         self.stop_token_id = stop_token_id
@@ -73,59 +91,47 @@ class CustomStoppingCriteria(StoppingCriteria):
 
 stopping_criteria = StoppingCriteriaList([CustomStoppingCriteria(tokenizer.eos_token_id)])
 
-
-# Step 6.5: Čisticí a kontrolní funkce
+# --------------------------
+# Step 6.5: Utility funkce
+# --------------------------
 def clean_response(text):
     """Odstraní HTML, Markdown a redundantní mezery."""
     original_text = text
-    text = re.sub(r"<[^>]+>", " ", text)   # odstraní HTML tagy
-    text = unescape(text)                  # dekóduje HTML entity
-    text = re.sub(r"[*#`_~]+", "", text)   # odstraní markdown znaky
+    text = re.sub(r"<[^>]+>", " ", text)
+    text = unescape(text)
+    text = re.sub(r"[*#`_~]+", "", text)
     text = re.sub(r"\s+", " ", text).strip()
     if text != original_text:
-        print("🧹 Response cleaned from HTML/Markdown artifacts.")
+        print("🧹 Cleaned response.")
     return text
 
 
 def remove_repetitions(text, similarity_threshold=0.8):
-    """
-    Pokud se opakují stejné věty (např. 'I'm TinyGPT...' 8x),
-    ponechá pouze první.
-    """
+    """Odstraní opakující se věty."""
     sentences = re.split(r'(?<=[.!?])\s+', text)
     if len(sentences) <= 1:
         return text
-
     unique_sentences = []
     for sent in sentences:
         sent_clean = sent.strip()
         if not sent_clean:
             continue
-        if not unique_sentences:
-            unique_sentences.append(sent_clean)
-            continue
-        ratio = SequenceMatcher(None, sent_clean, unique_sentences[-1]).ratio()
-        if ratio < similarity_threshold:
+        if not unique_sentences or SequenceMatcher(None, sent_clean, unique_sentences[-1]).ratio() < similarity_threshold:
             unique_sentences.append(sent_clean)
-
-    if len(unique_sentences) < len(sentences):
-        print("🧩 Repetitive content detected and reduced.")
     return " ".join(unique_sentences)
 
+
 def truncate_to_last_sentence(text):
     """Zkrátí text na poslední dokončenou větu."""
     sentences = re.split(r'(?<=[.!?])\s+', text)
-    if sentences and sentences[-1].strip():
-        # Najde poslední větu, která končí na . ? !
-        for i in range(len(sentences) - 1, -1, -1):
-            if re.search(r'[.!?]$', sentences[i].strip()):
-                return " ".join(sentences[:i+1]).strip()
-        # If no sentence ends with . ? !, return the whole text after cleaning
-        return text.strip()
+    for i in range(len(sentences) - 1, -1, -1):
+        if re.search(r'[.!?]$', sentences[i].strip()):
+            return " ".join(sentences[:i+1]).strip()
     return text.strip()
 
-
+# --------------------------
 # Step 7: Generování odpovědi
+# --------------------------
 def generate_response(
     user_input,
     max_length=2048,
@@ -139,9 +145,8 @@ def generate_response(
 ):
     try:
         prompt = f"{system_prompt}\n\nUser: {user_input}\nAssistant:"
-        device = "cuda:0" if torch.cuda.is_available() else "cpu"
         inputs = tokenizer(prompt, return_tensors="pt").to(device)
-        print(f"Input device: {inputs['input_ids'].device}")
+        print(f"📥 Input on device: {inputs['input_ids'].device}")
 
         with torch.no_grad():
             outputs = model.generate(
@@ -163,25 +168,23 @@ def generate_response(
         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         response = generated_text.split("Assistant:")[-1].strip()
 
-        # Vyčištění, odstranění opakování a zkrácení na poslední větu
         response = clean_response(response)
         response = remove_repetitions(response)
         response = truncate_to_last_sentence(response)
 
-
         return response
 
     except Exception as e:
-        print(f"Error during generation: {e}")
+        print(f"❌ Error during generation: {e}")
         return None
 
-# Step 8: Initialize Flask app
+# --------------------------
+# Step 8: Flask API
+# --------------------------
 app = Flask(__name__)
 
-# Step 9: Define API endpoint
 @app.route('/generate', methods=['POST'])
 def generate_text():
-    # Step 10: Get input, generate response, return JSON
     data = request.get_json()
     if not data or 'user_input' not in data:
         return jsonify({'error': 'Missing user_input parameter'}), 400
@@ -194,6 +197,8 @@ def generate_text():
 
     return jsonify({'response': generated_response})
 
-# Step 11: Run the Flask app
+# --------------------------
+# Step 9: Spuštění serveru
+# --------------------------
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860)