Update endpoints.py

endpoints.py

@@ -17,60 +17,85 @@ app.add_middleware(
     allow_headers=["*"],
     allow_credentials=True,
 )
-API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-v0.1"
-headers = {"Authorization": f"Bearer {key}"}
-
-def query(payload):
-	response = requests.post(API_URL, headers=headers, json=payload)
-	return response.json()
+# API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-v0.1"
+# headers = {"Authorization": f"Bearer {key}"}
 
+# def query(payload):
+# 	response = requests.post(API_URL, headers=headers, json=payload)
+# 	return response.json()
 
 
 
 
+def LLM(llm_name, length):
+    tokenizer = AutoTokenizer.from_pretrained(llm_name)
+    model = AutoModelForCausalLM.from_pretrained(llm_name)
+    pipe = pipeline("text-generation",
+                    model=model,
+                    tokenizer=tokenizer,
+                    max_length=length,
+                    do_sample=True,
+                    top_p=0.95,
+                    repetition_penalty=1.2,
+                   )
+    return pipe
 
 
 # tokenizer = AutoTokenizer.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 # base_model = AutoModelForCausalLM.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
-model = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(model)
-base_model = AutoModelForCausalLM.from_pretrained(model)
-pipe = pipeline("text-generation",
-                model=base_model,
-                tokenizer=tokenizer,
-                max_length=4000,
-                do_sample=True,
-                top_p=0.95,
-                repetition_penalty=1.2,
-               )
+# Mistral 7B
+mistral_llm = LLM("mistralai/Mistral-7B-v0.1",30000)
+
+
+# WizardCoder 13B
+wizard_llm = LLM("WizardLM/WizardCoder-Python-13B-V1.0",8000)
+
 # hf_llm = HuggingFacePipeline(pipeline=pipe)
 
+def ask_model(model, prompt):
+    if(model == 'mistral'):
+        return mistral_llm(prompt)
+    if(model == 'wizard'):
+        return wizard_llm(prompt)
+
+
 
 @app.get("/")
 def root():
     return {"message": "R&D LLM API"}
-@app.get("/get")
-def get():
-    result = pipe("name 5 programming languages",do_sample=False)
-    print(result)
-    return {"message": result}
+    
+# @app.get("/get")
+# def get():
+#     result = pipe("name 5 programming languages",do_sample=False)
+#     print(result)
+#     return {"message": result}
+
+
+
+
+
 
 
-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)
+async def ask_llm_endpoint(model:str, prompt: str):
+    result = ask_model(model,prompt)
     return {"result": result}
 
 
-@app.post("/ask_HFAPI")
-def ask_HFAPI_endpoint(prompt: str):
-    result = query(prompt)
-    return {"result": result}
+
+
+
+
+
+
+
+// APIs
+
+# @app.post("/ask_HFAPI")
+# def ask_HFAPI_endpoint(prompt: str):
+#     result = query(prompt)
+#     return {"result": result}
     
 from langchain.llms import OpenAI