Update app.py

app.py

@@ -1,19 +1,32 @@
 from fastapi import FastAPI
 import torch
-
+import os
+from llama_cpp import Llama
 from transformers import AutoModelForCausalLM, AutoTokenizer
 device = "cpu"
 
-tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
+access_token = os.getenv("access_token")
+
+tokenizer1 = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it", token=access_token)
+tokenizer3 = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
+
+llm1 = Llama.from_pretrained(
+    repo_id="Qwen/Qwen2-1.5B-Instruct-GGUF",
+    filename="*q8_0.gguf",
+    verbose=False
+)
 
-model = AutoModelForCausalLM.from_pretrained(
-    "Qwen/Qwen2-0.5B-Instruct",
-    device_map="auto"
+llm2 = Llama.from_pretrained(
+    repo_id="NexaAIDev/gemma-2-2b-it-GGUF",
+    filename="*q4_K_S.gguf",
+    verbose=False
 )
 
-model1 = AutoModelForCausalLM.from_pretrained(
-    "Qwen/Qwen2-1.5B-Instruct",
-    device_map="auto"
+llm3 = Llama.from_pretrained(
+    repo_id="microsoft/Phi-3-mini-4k-instruct-gguf",
+    filename="*q4.gguf",
+    verbose=False
 )
 
 app = FastAPI()
@@ -22,7 +35,7 @@ app = FastAPI()
 async def read_root():
     return {"Hello": "World!"}
 
-def modelResp(prompt):
+def modelResp1(prompt):
     messages = [
         {"role": "system", "content": "You are a helpful assistant, Sia, developed by Sushma. You will response in polity and brief."},
         {"role": "user", "content": "Who are you?"},
@@ -34,20 +47,16 @@ def modelResp(prompt):
         tokenize=False,
         add_generation_prompt=True
     )
-    model_inputs = tokenizer([text], return_tensors="pt").to(device)    
-    generated_ids = model.generate(
-        model_inputs.input_ids,
-        max_new_tokens=64,
-        do_sample=True
+    output = llm1(
+      text,
+      max_tokens=64,  # Generate up to 256 tokens
+      echo=False,  # Whether to echo the prompt
     )
-    generated_ids = [
-        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-    ] 
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    response = output['choices'][0]['text']
 
     return response
 
-def modelResp1(prompt):
+def modelResp2(prompt):
     messages = [
         {"role": "system", "content": "You are a helpful assistant, Sia, developed by Sushma. You will response in polity and brief."},
         {"role": "user", "content": "Who are you?"},
@@ -59,27 +68,50 @@ def modelResp1(prompt):
         tokenize=False,
         add_generation_prompt=True
     )
-    model_inputs = tokenizer([text], return_tensors="pt").to(device)    
-    generated_ids = model1.generate(
-        model_inputs.input_ids,
-        max_new_tokens=64,
-        do_sample=True
+    output = llm2(
+      text,
+      max_tokens=64,  # Generate up to 256 tokens
+      echo=False,  # Whether to echo the prompt
     )
-    generated_ids = [
-        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-    ] 
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    response = output['choices'][0]['text']
 
     return response
     
-@app.post("/modelapi")
+def modelResp3(prompt):
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant, Sia, developed by Sushma. You will response in polity and brief."},
+        {"role": "user", "content": "Who are you?"},
+        {"role": "assistant", "content": "I am Sia, a small language model created by Sushma."},
+        {"role": "user", "content": f"{prompt}"}
+    ]
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    output = llm2(
+      text,
+      max_tokens=64,  # Generate up to 256 tokens
+      echo=False,  # Whether to echo the prompt
+    )
+    response = output['choices'][0]['text']
+
+    return response
+    
+@app.post("/modelapi1")
 async def modelApi(data: dict):
     prompt = data.get("prompt")
-    response = modelResp(prompt)
+    response = modelResp1(prompt)
     return response
 
-@app.post("/modelapi1")
+@app.post("/modelapi2")
+async def modelApi(data: dict):
+    prompt = data.get("prompt")
+    response = modelResp2(prompt)
+    return response
+    
+@app.post("/modelapi3")
 async def modelApi1(data: dict):
     prompt = data.get("prompt")
-    response = modelResp1(prompt)
+    response = modelResp3(prompt)
     return response