Update app.py

app.py

@@ -56,13 +56,11 @@ def health():
 def chat():
     """Chat endpoint with BitNet streaming response"""
     global model_loaded, model, tokenizer
-
     if not model_loaded:
         return {
             "status": "initializing",
             "message": "Model is still loading. Please try again shortly."
         }, 503
-
     try:
         from transformers import TextIteratorStreamer
         data = request.get_json()
@@ -76,20 +74,16 @@ def chat():
         max_tokens = data.get("max_tokens", 512)
         temperature = data.get("temperature", 0.7)
         top_p = data.get("top_p", 0.95)
-
         messages = [{"role": "system", "content": system_message}]
         for user_msg, bot_msg in history:
             messages.append({"role": "user", "content": user_msg})
             messages.append({"role": "assistant", "content": bot_msg})
         messages.append({"role": "user", "content": message})
-
         prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
-
         streamer = TextIteratorStreamer(
             tokenizer, skip_prompt=True, skip_special_tokens=True
         )
-
         generate_kwargs = dict(
             **inputs,
             streamer=streamer,
@@ -98,17 +92,13 @@ def chat():
             top_p=top_p,
             do_sample=True,
         )
-
         thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
         thread.start()
-
         def generate():
             for new_text in streamer:
                 yield f"data: {json.dumps({'response': new_text})}\n\n"
             yield "data: [DONE]\n\n"
-
         return Response(generate(), mimetype="text/event-stream")
-
     except Exception as e:
         print("Error during chat:", e)
         return {"error": str(e)}, 500
@@ -117,24 +107,19 @@ def chat():
 def save_model():
     """Save model and tokenizer to Hugging Face Hub"""
     global model, tokenizer, model_loaded
-
     if not model_loaded:
         return {"error": "Model is still loading. Try again later."}, 503
-
     try:
         # Authenticate with Hugging Face
         token = request.json.get("token")
         if not token:
             return {"error": "Hugging Face token required"}, 400
         login(token=token)
-
         # Define repository
-        repo_id = "priyanshu/playwebit"
+        repo_id = "mike23415/playwebit"
         save_directory = "/tmp/playwebit"
-
         # Create temporary directory
         os.makedirs(save_directory, exist_ok=True)
-
         # Save custom model class (replace with actual implementation)
         custom_model_code = """
 from transformers import PreTrainedModel
@@ -154,15 +139,12 @@ class BitNetForCausalLM(PreTrainedModel):
 """
         with open(os.path.join(save_directory, "custom_bitnet.py"), "w") as f:
             f.write(custom_model_code)
-
         # Save configuration
         model.config.save_pretrained(save_directory)
-
         # Save model and tokenizer
         print("Saving model and tokenizer...")
         model.save_pretrained(save_directory, safe_serialization=True, max_shard_size="5GB")
         tokenizer.save_pretrained(save_directory)
-
         # Update config.json to reference custom class
         import json
         config_path = os.path.join(save_directory, "config.json")
@@ -171,7 +153,6 @@ class BitNetForCausalLM(PreTrainedModel):
         config_json["architectures"] = ["BitNetForCausalLM"]
         with open(config_path, "w") as f:
             json.dump(config_json, f, indent=2)
-
         # Try TensorFlow conversion
         try:
             from transformers import TFAutoModelForCausalLM
@@ -180,23 +161,4 @@ class BitNetForCausalLM(PreTrainedModel):
             tf_model.save_pretrained(save_directory)
             print("TensorFlow weights saved.")
         except Exception as e:
-            print(f"Error converting to TensorFlow: {e}")
-
-        # Upload to Hugging Face Hub
-        api = HfApi()
-        print(f"Uploading to {repo_id}...")
-        api.upload_folder(
-            folder_path=save_directory,
-            repo_id=repo_id,
-            repo_type="model",
-            commit_message="Upload PlayWeBit model, tokenizer, and custom class"
-        )
-
-        return {"message": f"Model uploaded to https://huggingface.co/{repo_id}"}
-
-    except Exception as e:
-        print("Error saving model:", e)
-        return {"error": str(e)}, 500
-
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
+            print(f"Error converting to TensorFlow: {e}")