debug

data/models/llama3-1-70b.py

@@ -1,25 +1,19 @@
 from flask import Flask, request, jsonify
-from huggingface_hub import login
+from huggingface_hub import login, InferenceClient
 import spaces
-import transformers
-import torch
 import os
 
+
+
 # Initialize Flask app
 app = Flask(__name__)
 
 api_key = os.getenv("UCODE_SECRET")
 login(api_key,add_to_git_credential=True)
 
-model_id = "meta-llama/Meta-Llama-3.1-70B-Instruct"
+client = InferenceClient()
 
-pipeline = transformers.pipeline(
-    "text-generation",
-    model=model_id,
-    model_kwargs={"torch_dtype": torch.bfloat16,"quantization_config": {"load_in_4bit": True}},
-    device="cuda",
-    token=True
-)
+model_id = "meta-llama/Meta-Llama-3.1-70B-Instruct"
 
 @app.route('/chat', methods=['POST'])
 @spaces.GPU(enable_queue=True)
@@ -31,6 +25,15 @@ def chat_completion():
     temperature = data[0].get('temperature', 0.7)
     top_p = data[0].get('top_p', 0.95)
 
+    chat = client.chat.completions.create(
+        model=model_id,
+        messages=user_input,
+        stream=False,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p
+    )
+
     try:
         outputs = pipeline(
             user_input,