Create app.py

app.py

@@ -0,0 +1,72 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import requests
+from bs4 import BeautifulSoup
+
+app = FastAPI()
+
+MODEL_NAME = "microsoft/phi-2"
+
+print("Loading Phi-2...")
+
+torch.set_num_threads(2)
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+model.to("cpu")
+
+print("Model loaded!")
+
+# -------- REQUEST SCHEMA --------
+class RequestData(BaseModel):
+    prompt: str
+    use_search: bool = False
+
+
+# -------- WEB SEARCH FUNCTION --------
+def search_web(query):
+    url = f"https://duckduckgo.com/html/?q={query}"
+    headers = {"User-Agent": "Mozilla/5.0"}
+
+    response = requests.get(url, headers=headers)
+    soup = BeautifulSoup(response.text, "html.parser")
+
+    results = []
+    for a in soup.select("a.result__a"):
+        results.append(a.get_text())
+
+    return " ".join(results[:5])
+
+
+# -------- GENERATE FUNCTION --------
+def generate_text(prompt):
+    formatted = f"Instruct: {prompt}\nOutput:"
+
+    inputs = tokenizer(formatted, return_tensors="pt")
+
+    outputs = model.generate(
+        inputs["input_ids"],
+        max_new_tokens=60,
+        temperature=0.7
+    )
+
+    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return result.split("Output:")[-1].strip()
+
+
+# -------- API ENDPOINT --------
+@app.post("/generate")
+def generate(data: RequestData):
+    prompt = data.prompt
+
+    if data.use_search:
+        web_data = search_web(prompt)
+        prompt = f"{prompt}\n\nWeb Info: {web_data}"
+
+    response = generate_text(prompt)
+
+    return {
+        "response": response
+    }