Upload main.py

main.py

@@ -0,0 +1,21 @@
+from ctransformers import AutoModelForCausalLM
+from fastapi import FastAPI
+from pydantic import BaseModel
+
+# Load Gemma-2B model
+model = AutoModelForCausalLM.from_pretrained("gemma-2b.Q4_K_S.gguf", max_new_tokens=8192, threads=3)
+
+# Define Pydantic object for input validation
+class Validation(BaseModel):
+    prompt: str
+
+# Initialize FastAPI
+app = FastAPI()
+
+# Define route for generating text
+@app.post("/generate_text")
+async def generate_text(item: Validation):
+    # Construct the input prompt
+    input_prompt = f"<|im_start|>system\n{{Below is an instruction that describes a task. Write a response that appropriately completes the request.}}<|im_end|>\n<|im_start|>user\n{{{item.prompt}}}<|im_end|>\n<|im_start|>assistant"
+
+    return llm(input_prompt)