misalsathsara commited on
Commit
234a50d
·
verified ·
1 Parent(s): 973ab13

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +71 -0
app.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI
2
+ from pydantic import BaseModel
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM
4
+ import torch
5
+ import re
6
+
7
+ app = FastAPI()
8
+
9
+ # Load model and tokenizer from Hugging Face
10
+ model_id = "misalsathsara/phi1.5-js-codegen"
11
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
12
+ model = AutoModelForCausalLM.from_pretrained(model_id)
13
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
14
+ model.to(device)
15
+ model.eval()
16
+
17
+ # Your system prompt
18
+ system_prompt = """
19
+ You are a smart javascript assistant that only generates only the best simple javascript functions without any comments like this:
20
+ function transform(row) {
21
+ row['Latitude'] = row['Location'].split(',')[0];
22
+ row['Longitude'] = row['Location'].split(',')[1];
23
+ return row;
24
+ }
25
+ when user gives a prompt like "convert the location field into separate latitude and longitude fields".
26
+ Generate simple javascript functions that should take a single row of data as input and the generated function name is always transform.
27
+ The user may use the words column, item or field to mean each column.
28
+ Guard against null and undefined for items in the row.
29
+ ${fieldList}
30
+ Field names are case sensitive.
31
+ For parsing something into a date, assume a function called parseAnyDate is available.
32
+ If the code requires some numeric calculation - ensure the value is converted to a number first. Don't assume its always the correct data type.
33
+ When doing any string comparison, make it case insensitive.
34
+ When replacing characters in a string, make sure to use the correct replacement literal. For example, to replace hyphens with spaces, use: .replace(/-/g, ' ')
35
+ The function should not include a single comment before or after the function.
36
+ Don't add any text except for the function code.
37
+ Don't add any markdown block markers either.
38
+ Every function must end with return row;
39
+ """
40
+
41
+ # Define the expected request body
42
+ class RequestData(BaseModel):
43
+ instruction: str
44
+
45
+ # POST endpoint
46
+ @app.post("/generate")
47
+ def generate_code(data: RequestData):
48
+ instruction = data.instruction
49
+ full_prompt = system_prompt + f"\n### Instruction:\n{instruction}\n\n### Response:\n"
50
+
51
+ # Tokenize input
52
+ input_ids = tokenizer(full_prompt, return_tensors="pt").input_ids.to(device)
53
+
54
+ with torch.no_grad():
55
+ output_ids = model.generate(
56
+ input_ids,
57
+ max_new_tokens=200,
58
+ temperature=0.3,
59
+ top_k=50,
60
+ top_p=0.95,
61
+ do_sample=True,
62
+ pad_token_id=tokenizer.eos_token_id
63
+ )
64
+
65
+ generated_text = tokenizer.decode(output_ids[0][input_ids.shape[-1]:], skip_special_tokens=True)
66
+
67
+ # Extract clean JS function
68
+ match = re.search(r"function\s*\(.*?\)\s*{.*?return row;\s*}", generated_text, re.DOTALL)
69
+ clean_output = match.group(0).strip() if match else generated_text.strip()
70
+
71
+ return {"result": clean_output}