Update main.py

main.py

@@ -26,7 +26,26 @@ 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)
+
+    device = "cpu"
+    
+    prompt = item.prompt
+
+    encodeds = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
+
+    model_inputs = encodeds.to(device)
+
+
+    generated_ids = model.generate(**model_inputs, item.max_new_tokens, do_sample=True, pad_token_id=tokenizer.eos_token_id, temperature=temperature)
+    decoded = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    return decoded
+
 
 @app.post("/generate/")
 async def generate_text(item: Item):
-    return {"response": item.prompt}
+    return {"response": generate(item)}