LLM micro service extract input model name (opea-project#447)

No input model from UI. vLLM and OLLAMA get the model name from the envrionment TGI mode model form ModelID Signed-off-by: Wang, Xigui <[email protected]> Signed-off-by: BaoHuiling <[email protected]>
siddhivelankar23 · Aug 15, 2024 · cc5b741 · cc5b741
1 parent a497345
commit cc5b741
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/comps/llms/text-generation/ollama/llm.py b/comps/llms/text-generation/ollama/llm.py
@@ -21,7 +21,7 @@
 def llm_generate(input: LLMParamsDoc):
     ollama = Ollama(
         base_url=ollama_endpoint,
-        model=input.model,
+        model=input.model if input.model else model_name,
         num_predict=input.max_new_tokens,
         top_k=input.top_k,
         top_p=input.top_p,
@@ -49,4 +49,5 @@ async def stream_generator():
 
 if __name__ == "__main__":
     ollama_endpoint = os.getenv("OLLAMA_ENDPOINT", "http://localhost:11434")
+    model_name = os.getenv("OLLAMA_MODEL", "meta-llama/Meta-Llama-3-8B-Instruct")
     opea_microservices["opea_service@llm_ollama"].start()