Update app.py

app.py

@@ -1,7 +1,7 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from huggingface_hub import hf_hub_download
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from llama_cpp import Llama
 
 # Definição do modelo de dados de entrada
 class Question(BaseModel):
@@ -16,12 +16,14 @@ model_basename = "unsloth.Q8_0.gguf"
 model_path = hf_hub_download(repo_id=model_name_or_path, filename=model_basename)
 print(f"Model path: {model_path}")
 
-# Carregar o tokenizador e o modelo
-tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
-model = AutoModelForCausalLM.from_pretrained(model_name_or_path).to("cuda")
-
-# Ativar inferência otimizada
-FastLanguageModel.for_inference(model)  # Ativar inferência otimizada
+# Configuração do modelo com llama_cpp
+lcpp_llm = Llama(
+    model_path=model_path,
+    n_threads=2,
+    n_batch=512,
+    n_gpu_layers=-1,
+    n_ctx=4096,
+)
 
 # Formato de prompt utilizado no fine-tuning
 alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
@@ -42,14 +44,18 @@ def get_response(text: str) -> str:
         text,
         ""
     )
-    inputs = tokenizer([formatted_prompt], return_tensors="pt").to("cuda")
-    
-    # Gerar resposta
-    outputs = model.generate(**inputs, max_new_tokens=64, use_cache=True)
-    response = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
-    
-    print(f"Generated Response: {response}")
-    return response
+    response = lcpp_llm(
+        prompt=formatted_prompt,
+        max_tokens=256,
+        temperature=0.5,
+        top_p=0.95,
+        top_k=50,
+        stop=['\n'],  # Tente usar um token de parada específico
+        echo=True
+    )
+    response_text = response['choices'][0]['text']
+    print(f"Generated Response: {response_text}")
+    return response_text
 
 # Endpoint para receber uma questão e retornar a resposta
 @app.post("/ask")
@@ -62,3 +68,4 @@ if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)
 
+