Update endpoints.py

endpoints.py

@@ -26,8 +26,8 @@ app.add_middleware(
 # 	response = requests.post(API_URL, headers=headers, json=payload)
 # 	return response.json()
 
-from fastapi import FastAPI, HTTPException, Body
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
+# from fastapi import FastAPI, HTTPException, Body
+# from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
 
 # model = GPT2LMHeadModel.from_pretrained("EleutherAI/gpt-neo-2.7B")
@@ -37,46 +37,46 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer
 tokenizer = AutoTokenizer.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 model = AutoModelForCausalLM.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 
-# pipe = pipeline("text-generation",
-#                 model=base_model,
-#                 tokenizer=tokenizer,
-#                 max_length=4000,
-#                 do_sample=True,
-#                 top_p=0.95,
-#                 repetition_penalty=1.2,
-#                )
-# hf_llm = HuggingFacePipeline(pipeline=pipe)
+pipe = pipeline("text-generation",
+                model=base_model,
+                tokenizer=tokenizer,
+                max_length=4000,
+                do_sample=True,
+                top_p=0.95,
+                repetition_penalty=1.2,
+               )
+hf_llm = HuggingFacePipeline(pipeline=pipe)
 
 
-class ChatRequest(BaseModel):
-    messages: list
-    temperature: float = 1.0
-    max_tokens: int = 50
-    stream: bool = False
+# class ChatRequest(BaseModel):
+#     messages: list
+#     temperature: float = 1.0
+#     max_tokens: int = 50
+#     stream: bool = False
 
-class ChatResponse(BaseModel):
-    response: str
+# class ChatResponse(BaseModel):
+#     response: str
 
-@app.post("/v1/chat/completions", response_model=ChatResponse)
-async def chat_completions(request: ChatRequest):
-    try:
-        # Prepare input prompt
-        input_prompt = ""
-        for message in request.messages:
-            role = message.get('role', 'user')
-            content = message.get('content', '')
-            input_prompt += f"{role}: {content}\n"
+# @app.post("/v1/chat/completions", response_model=ChatResponse)
+# async def chat_completions(request: ChatRequest):
+#     try:
+#         # Prepare input prompt
+#         input_prompt = ""
+#         for message in request.messages:
+#             role = message.get('role', 'user')
+#             content = message.get('content', '')
+#             input_prompt += f"{role}: {content}\n"
 
-        # Tokenize and generate response
-        input_ids = tokenizer.encode(input_prompt, return_tensors='pt')
-        output = model.generate(input_ids, max_length=1024, temperature=request.temperature, max_tokens=request.max_tokens)
+#         # Tokenize and generate response
+#         input_ids = tokenizer.encode(input_prompt, return_tensors='pt')
+#         output = model.generate(input_ids, max_length=1024, temperature=request.temperature, max_tokens=request.max_tokens)
 
-        # Decode and send response
-        response = tokenizer.decode(output[0], skip_special_tokens=True)
-        return {"response": response}
+#         # Decode and send response
+#         response = tokenizer.decode(output[0], skip_special_tokens=True)
+#         return {"response": response}
 
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+#     except Exception as e:
+#         raise HTTPException(status_code=500, detail=str(e))
 
 
 @app.get("/")
@@ -89,15 +89,15 @@ def root():
 #     return {"message": result}
 
 
-# async def askLLM(prompt):
-#     output = pipe(prompt,do_sample=False)
-#     return output
+async def askLLM(prompt):
+    output = pipe(prompt,do_sample=False)
+    return output
 
-# @app.post("/ask_llm")
-# async def ask_llm_endpoint(prompt: str):
-#     # result = await askLLM(prompt)
-#     result = pipe(prompt,do_sample=False)
-#     return {"result": result}
+@app.post("/ask_llm")
+async def ask_llm_endpoint(prompt: str):
+    # result = await askLLM(prompt)
+    result = pipe(prompt,do_sample=False)
+    return {"result": result}
 
 
 # @app.post("/ask_HFAPI")