Update app.py

app.py

@@ -1,13 +1,18 @@
 from fastapi import FastAPI, Request, HTTPException
-import requests
 import os
 from typing import Dict, Any
+from transformers import pipeline
 
 app = FastAPI()
 
+# Initialisation du pipeline pour l'API d'inférence Hugging Face
 BOT_USERNAME = "@DiscussionBot"
-INFERENCE_URL = "api-inference.huggingface.co/models/facebook/bart-large-cnn"
-PROMPT = "Faites comme si vous étiez un robot qui répond aux discussions sur l'apprentissage automatique et répondez au commentaire suivant :\n"
+MODEL = "meta-llama/Llama-3.2-1B-Instruct"
+try:
+    pipe = pipeline("text-generation", model=MODEL, token=os.getenv("HF_TOKEN"))
+except Exception as e:
+    print(f"Failed to initialize pipeline: {str(e)}")
+    raise
 
 @app.get("/")
 async def root(request: Request) -> Dict[str, Any]:
@@ -45,39 +50,41 @@ async def webhook(request: Request) -> Dict[str, Any]:
         ):
             print(f"Processing comment: {data['comment']['content']}")
 
-            response = requests.post(
-                INFERENCE_URL,
-                headers={"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"},
-                json={"inputs": PROMPT + data["comment"]["content"]},
-            )
+            # Préparation du prompt au format "messages"
+            messages = [
+                {
+                    "role": "user",
+                    "content": f"Faites comme si vous étiez un robot qui répond aux discussions sur l'apprentissage automatique et répondez au commentaire suivant :\n{data['comment']['content']}"
+                }
+            ]
 
-            if response.status_code == 200:
-                output = response.json()
-                print(f"BLOOM response: {output}")
-                continuation_text = output[0]["generated_text"].replace(
-                    PROMPT + data["comment"]["content"], ""
-                ).strip()
+            # Inférence avec le pipeline
+            try:
+                output = pipe(messages, max_new_tokens=100, num_return_sequences=1)
+                print(f"Pipeline response: {output}")
+                continuation_text = output[0]["generated_text"][-1]["content"].strip()
+            except Exception as e:
+                error_message = f"Inference failed: {str(e)}"
+                print(error_message)
+                raise HTTPException(status_code=500, detail=error_message)
 
-                comment_url = data["discussion"]["url"]["api"] + "/comment"
-                print(f"Posting to comment URL: {comment_url}")
-                comment_response = requests.post(
-                    comment_url,
-                    headers={
-                        "Authorization": f"Bearer {os.getenv('HF_TOKEN')}",
-                        "Content-Type": "application/json",
-                    },
-                    json={"comment": continuation_text},
-                )
+            # Publication du commentaire
+            comment_url = data["discussion"]["url"]["api"] + "/comment"
+            print(f"Posting to comment URL: {comment_url}")
+            comment_response = requests.post(
+                comment_url,
+                headers={
+                    "Authorization": f"Bearer {os.getenv('HF_TOKEN')}",
+                    "Content-Type": "application/json",
+                },
+                json={"comment": continuation_text},
+            )
 
-                if comment_response.status_code == 200:
-                    print("Comment posted successfully")
-                    return {"success": True}
-                else:
-                    error_message = f"Failed to post comment: {comment_response.text}"
-                    print(error_message)
-                    raise HTTPException(status_code=500, detail=error_message)
+            if comment_response.status_code == 200:
+                print("Comment posted successfully")
+                return {"success": True}
             else:
-                error_message = f"Inference API failed: {response.text}"
+                error_message = f"Failed to post comment: {comment_response.text}"
                 print(error_message)
                 raise HTTPException(status_code=500, detail=error_message)