Update app.py

app.py

@@ -1,12 +1,13 @@
 import os
-import requests as req
 from fastapi import FastAPI
 from pydantic import BaseModel
+from huggingface_hub import InferenceClient
 
 app = FastAPI()
-
-HF_TOKEN = os.environ.get("HF_TOKEN")
-API_URL = "https://api-inference.huggingface.co/models/google/flan-t5-large"
+client = InferenceClient(
+    provider="hf-inference",
+    api_key=os.environ.get("HF_TOKEN")
+)
 
 class Request(BaseModel):
     prompt: str
@@ -17,15 +18,10 @@ def home():
 
 @app.post("/generate")
 def generate_text(request: Request):
-    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-    payload = {"inputs": request.prompt}
-    response = req.post(API_URL, headers=headers, json=payload)
-    result = response.json()
-    
-    # Handle different response formats
-    if isinstance(result, list) and len(result) > 0:
-        return {"response": result[0].get("generated_text", str(result))}
-    elif isinstance(result, dict):
-        return {"response": result.get("generated_text", str(result))}
-    else:
-        return {"response": str(result)}
+    messages = [{"role": "user", "content": request.prompt}]
+    completion = client.chat.completions.create(
+        model="Qwen/Qwen2.5-72B-Instruct",
+        messages=messages,
+        max_tokens=500
+    )
+    return {"response": completion.choices[0].message.content}