Update endpoints.py

endpoints.py

@@ -6,17 +6,6 @@ import os
 import requests
 # from langchain.llms.huggingface_pipeline import HuggingFacePipeline
 
-key = os.environ.get("huggingface_key")
-openai_api_key = os.environ.get("openai_key")
-app = FastAPI(openapi_url="/api/v1/LLM/openapi.json", docs_url="/api/v1/LLM/docs")
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-    allow_credentials=True,
-)
 # API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-v0.1"
 # headers = {"Authorization": f"Bearer {key}"}
 
@@ -24,33 +13,34 @@ app.add_middleware(
 # 	response = requests.post(API_URL, headers=headers, json=payload)
 # 	return response.json()
 
-
-
-
-def LLM(llm_name, length):
-    print(llm_name)
-    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
-
+# def LLM(llm_name, length):
+#     print(llm_name)
+#     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
+# Load model directly
+# Use a pipeline as a high-level helper
+from transformers import pipeline
+
+pipe = pipeline("text-generation", model="mistralai/Mistral-7B-v0.1")
 
 # tokenizer = AutoTokenizer.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 # base_model = AutoModelForCausalLM.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 # Mistral 7B
-mistral_llm = LLM("mistralai/Mistral-7B-v0.1",30000)
-
+# mistral_llm = LLM("mistralai/Mistral-7B-v0.1",30000)
+mistral_llm = pipe
 
 # WizardCoder 13B
-wizard_llm = LLM("WizardLM/WizardCoder-Python-13B-V1.0",8000)
-
+# wizard_llm = LLM("WizardLM/WizardCoder-Python-13B-V1.0",8000)
+wizard_llm = pipe
 # hf_llm = HuggingFacePipeline(pipeline=pipe)
 
 def ask_model(model, prompt):
@@ -61,6 +51,22 @@ def ask_model(model, prompt):
 
 
 
+
+
+
+key = os.environ.get("huggingface_key")
+openai_api_key = os.environ.get("openai_key")
+app = FastAPI(openapi_url="/api/v1/LLM/openapi.json", docs_url="/api/v1/LLM/docs")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+    allow_credentials=True,
+)
+
+
 @app.get("/")
 def root():
     return {"message": "R&D LLM API"}