Update main.py

main.py

@@ -6,8 +6,6 @@ import uvicorn
 
 app = FastAPI()
 
-client = InferenceClient("google/gemma-1.1-7b-it")
-
 class Item(BaseModel):
     prompt: str
     system_prompt: str
@@ -16,28 +14,7 @@ class Item(BaseModel):
     top_p: float = 0.15
     repetition_penalty: float = 1.0
 
-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,
-    )
-
-    stream = client.text_generation(item.prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
-
-    for response in stream:
-        output += response.token.text
-    return output
 
 @app.post("/generate/")
 async def generate_text(item: Item):
-    return {"response": generate(item)}
+    return {"response": item.prompt}