add generation code

app.py

@@ -1,4 +1,16 @@
 from fastapi import FastAPI
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+import torch
+from pydantic import BaseModel, Field
+
+class Request(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,
 
 app = FastAPI()
 
@@ -6,6 +18,25 @@ app = FastAPI()
 def greet_json():
     return {"Hello": "World!"}
 
-@app.post("/")
-def greet_json():
-    return {"Hello": "World!"}
+@app.post("/generate")
+def generate(req: Request):
+    model_name_or_id = "AI4Chem/ChemLLM-7B-Chat"
+
+    model = AutoModelForCausalLM.from_pretrained(model_name_or_id,trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_id,trust_remote_code=True)
+
+    inputs = tokenizer(req.prompt, return_tensors="pt")
+
+    generation_config = GenerationConfig(
+        do_sample=req.do_sample,
+        top_k=req.top_k,
+        temperature=req.temperature,
+        max_new_tokens=req.max_new_tokens,
+        repetition_penalty=req.repetition_penalty,
+        pad_token_id=tokenizer.eos_token_id
+    )
+
+    outputs = model.generate(**inputs, generation_config=generation_config)
+    # print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+
+    return {"text": tokenizer.decode(outputs[0], skip_special_tokens=True)}

requirements.txt

@@ -1,2 +1,4 @@
 fastapi
 uvicorn[standard]
+transformers
+torch