fix generation code

app.py

@@ -3,14 +3,13 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 import torch
 from pydantic import BaseModel, Field
 
-class Request(BaseModel):
+class RequestGenerate(BaseModel):
     prompt: str 
-    response: str = Field(description="")
-    do_sample: bool=True,
-    top_k: int =1,
-    temperature: float=0.9,
-    max_new_tokens: int=500,
-    repetition_penalty: float=1.5,
+    do_sample: bool = Field(default=True, example=True)
+    top_k: int = Field(default=1, example=1),
+    temperature: float = Field(default=0.9, example=0.9),
+    max_new_tokens: int = Field(default=500, example=500),
+    repetition_penalty: float = Field(default=1.5, example=1.5),
 
 app = FastAPI()
 
@@ -19,8 +18,9 @@ def greet_json():
     return {"Hello": "World!"}
 
 @app.post("/generate")
-def generate(req: Request):
+def generate(req: RequestGenerate):
     model_name_or_id = "AI4Chem/ChemLLM-7B-Chat"
+    # model_name_or_id = "AI4Chem/CHEMLLM-2b-1_5"
 
     model = AutoModelForCausalLM.from_pretrained(model_name_or_id,trust_remote_code=True)
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_id,trust_remote_code=True)

requirements.txt

@@ -2,3 +2,5 @@ fastapi
 uvicorn[standard]
 transformers
 torch
+einops
+sentencepiece