Update app.py

app.py

@@ -1,10 +1,9 @@
-import gradio as gr
-from text_generation import InferenceAPIClient
 import os
+import gradio as gr
+from huggingface_hub import InferenceClient
 
-# ✅ Choose your model (you can change this to another instruct model)
-# client = InferenceAPIClient("deepseek-ai/DeepSeek-R1-0528", token=os.getenv("HF_TOKEN"))
-client = InferenceAPIClient("bigscience/bloomz")
+# ✅ Set up Hugging Face client with token from environment
+llm = InferenceClient(token=os.getenv("HF_TOKEN"))
 
 # 🧠 Function to generate software architecture
 def generate_software_spec(name, description, architecture, components, deployment, platform, extra):
@@ -49,17 +48,16 @@ Return the result in this format:
 ```
 """
 
-    # Token Streaming
-    text = ""
-    for response in client.generate_stream(prompt):
-        if not response.token.special:
-            text += response.token.text
-    return text
-    # Generate from model (stream=True recommended for large output)
-    #result = ""
-    #for chunk in client.text_generation(prompt, stream=True, max_new_tokens=1024, temperature=0.7, stop=["</s>"]):
-    #    result += chunk.token.text
-    #return result
+    # Generate response from Hugging Face Inference Client
+    response = llm.chat_completion(
+        messages=[
+            {"role": "system", "content": "You are a helpful software architecture assistant."},
+            {"role": "user", "content": prompt}
+        ],
+        model="deepseek-ai/DeepSeek-R1-0528",
+        max_tokens=2048,
+    )
+    return response.choices[0].message["content"] if response and response.choices else "Error: No response received."
 
 # 🎨 Gradio UI
 with gr.Blocks() as demo: