Update main.py

main.py

@@ -1,11 +1,10 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from huggingface_hub import InferenceClient
-import uvicorn
-
+from huggingface_hub import InferenceClient, BadRequestError
 
 app = FastAPI()
 
+# Use your model
 client = InferenceClient("ManojINaik/codsw")
 
 class Item(BaseModel):
@@ -26,29 +25,37 @@ def format_prompt(message, history):
     return prompt
 
 def generate(item: Item):
-    temperature = float(item.temperature)
-    if temperature < 1e-2:
-        temperature = 1e-2
-    top_p = float(item.top_p)
-
-    generate_kwargs = dict(
-        temperature=temperature,
-        max_new_tokens=item.max_new_tokens,
-        top_p=top_p,
-        repetition_penalty=item.repetition_penalty,
-        do_sample=True,
-        seed=42,
-    )
-
-    formatted_prompt = format_prompt(f"{item.system_prompt}, {item.prompt}", item.history)
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
-
-    for response in stream:
-        output += response.token.text
-    return output
+    try:
+        # Ensure valid temperature
+        temperature = max(float(item.temperature), 1e-2)
+        top_p = float(item.top_p)
+
+        generate_kwargs = {
+            "temperature": temperature,
+            "max_new_tokens": item.max_new_tokens,
+            "top_p": top_p,
+            "repetition_penalty": item.repetition_penalty,
+            "do_sample": True,
+            "seed": 42,
+        }
+
+        # Format the prompt
+        formatted_prompt = format_prompt(f"{item.system_prompt}, {item.prompt}", item.history)
+
+        # Call text_generation on your model
+        stream = client.text_generation(
+            inputs=formatted_prompt,
+            **generate_kwargs,
+            stream=True,
+        )
+        output = "".join([response.token.text for response in stream])
+        return output
+
+    except BadRequestError as e:
+        raise HTTPException(status_code=400, detail=f"Bad request: {str(e)}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Internal error: {str(e)}")
 
 @app.post("/generate/")
 async def generate_text(item: Item):
     return {"response": generate(item)}
-